From 0e333252d606ce4854512160f403b1df1941d17b Mon Sep 17 00:00:00 2001 From: ninaburg <83002751+ninaburg@users.noreply.github.com> Date: Thu, 28 Sep 2023 17:01:40 +0200 Subject: [PATCH] Porting the tracer advection to gt4py (#252) Merge advection to main --- .github/workflows/icon4py-qa.yml | 4 +- model/README.md | 2 + model/atmosphere/advection/.flake8 | 42 + .../advection/.pre-commit-config.yaml | 114 +++ model/atmosphere/advection/README.md | 9 + .../advection/advection_tests/__init__.py | 12 + .../advection/advection_tests/conftest.py | 33 + .../test_btraj_dreg_stencil_01.py | 63 ++ .../test_btraj_dreg_stencil_02.py | 71 ++ .../test_btraj_dreg_stencil_03.py | 228 +++++ .../test_divide_flux_area_list_stencil_01.py | 799 ++++++++++++++++++ .../test_divide_flux_area_list_stencil_02.py | 333 ++++++++ .../test_face_val_ppm_stencil_01.py | 93 ++ .../test_face_val_ppm_stencil_02.py | 88 ++ .../test_face_val_ppm_stencil_02a.py | 54 ++ .../test_face_val_ppm_stencil_02b.py | 46 + .../test_face_val_ppm_stencil_02c.py | 48 ++ .../test_face_val_ppm_stencil_05.py | 86 ++ .../test_hflux_ffsl_hybrid_stencil_01a.py | 205 +++++ .../test_hflux_ffsl_hybrid_stencil_02.py | 54 ++ .../test_hflx_limiter_mo_stencil_01a.py | 69 ++ .../test_hflx_limiter_mo_stencil_01b.py | 122 +++ .../test_hflx_limiter_mo_stencil_02.py | 129 +++ .../test_hflx_limiter_mo_stencil_03.py | 122 +++ .../test_hflx_limiter_mo_stencil_04.py | 59 ++ .../test_hflx_limiter_pd_stencil_01.py | 76 ++ .../test_hflx_limiter_pd_stencil_02.py | 125 +++ .../test_hor_adv_stencil_01.py | 80 ++ ..._mo_advection_traj_btraj_compute_o1_dsl.py | 6 +- ...st_prep_gauss_quadrature_c_list_stencil.py | 505 +++++++++++ .../test_prep_gauss_quadrature_c_stencil.py | 467 ++++++++++ .../test_rbf_intp_edge_stencil_01.py | 54 ++ .../test_recon_lsq_cell_c_svd_stencil.py | 340 ++++++++ .../test_recon_lsq_cell_l_svd_stencil.py | 77 ++ .../advection_tests/test_set_zero_c.py | 27 + .../advection_tests/test_set_zero_c_k.py | 27 + .../test_step_advection_stencil_01.py | 62 ++ .../test_step_advection_stencil_02.py | 61 ++ .../test_step_advection_stencil_03.py | 55 ++ .../test_step_advection_stencil_04.py | 51 ++ .../test_upwind_hflux_miura3_stencil_01.py | 217 +++++ ...test_upwind_hflux_miura_cycl_stencil_01.py | 98 +++ ...test_upwind_hflux_miura_cycl_stencil_02.py | 101 +++ ...est_upwind_hflux_miura_cycl_stencil_03a.py | 49 ++ ...est_upwind_hflux_miura_cycl_stencil_03b.py | 53 ++ .../test_upwind_hflux_miura_stencil_01.py | 99 +++ .../test_upwind_vflux_ppm_stencil_01.py | 47 ++ .../test_vert_adv_stencil_01.py | 73 ++ .../test_vlimit_prbl_sm_stencil_01.py | 54 ++ .../test_vlimit_prbl_sm_stencil_02.py | 70 ++ model/atmosphere/advection/pyproject.toml | 120 +++ .../atmosphere/advection/requirements-dev.txt | 3 + model/atmosphere/advection/requirements.txt | 3 + .../model/atmosphere/advection/__init__.py | 33 + .../advection/btraj_dreg_stencil_01.py | 41 + .../advection/btraj_dreg_stencil_02.py | 46 + .../advection/btraj_dreg_stencil_03.py | 169 ++++ .../divide_flux_area_list_stencil_01.py | 744 ++++++++++++++++ .../divide_flux_area_list_stencil_02.py | 267 ++++++ .../advection/face_val_ppm_stencil_01.py | 90 ++ .../advection/face_val_ppm_stencil_02.py | 100 +++ .../advection/face_val_ppm_stencil_02a.py | 43 + .../advection/face_val_ppm_stencil_02b.py | 37 + .../advection/face_val_ppm_stencil_02c.py | 37 + .../advection/face_val_ppm_stencil_05.py | 67 ++ .../hflux_ffsl_hybrid_stencil_01a.py | 160 ++++ .../advection/hflux_ffsl_hybrid_stencil_02.py | 43 + .../advection/hflx_limiter_mo_stencil_01a.py | 50 ++ .../advection/hflx_limiter_mo_stencil_01b.py | 95 +++ .../advection/hflx_limiter_mo_stencil_02.py | 70 ++ .../advection/hflx_limiter_mo_stencil_03.py | 114 +++ .../advection/hflx_limiter_mo_stencil_04.py | 44 + .../advection/hflx_limiter_mo_stencil_05.py | 49 ++ .../advection/hflx_limiter_pd_stencil_01.py | 58 ++ .../advection/hflx_limiter_pd_stencil_02.py | 52 ++ .../advection/hor_adv_stencil_01.py | 60 ++ .../mo_advection_traj_btraj_compute_o1_dsl.py | 13 +- .../prep_gauss_quadrature_c_list_stencil.py | 435 ++++++++++ .../prep_gauss_quadrature_c_stencil.py | 396 +++++++++ .../model/atmosphere/advection/py.typed | 0 .../advection/rbf_intp_edge_stencil_01.py | 35 + .../advection/recon_lsq_cell_c_svd_stencil.py | 244 ++++++ .../advection/recon_lsq_cell_l_svd_stencil.py | 58 ++ .../model/atmosphere/advection/set_zero_c.py | 28 + .../atmosphere/advection/set_zero_c_k.py | 28 + .../advection/step_advection_stencil_01.py | 50 ++ .../advection/step_advection_stencil_02.py | 50 ++ .../advection/step_advection_stencil_03.py | 37 + .../advection/step_advection_stencil_04.py | 37 + .../upwind_hflux_miura3_stencil_01.py | 170 ++++ .../upwind_hflux_miura_cycl_stencil_01.py | 75 ++ .../upwind_hflux_miura_cycl_stencil_02.py | 77 ++ .../upwind_hflux_miura_cycl_stencil_03a.py | 36 + .../upwind_hflux_miura_cycl_stencil_03b.py | 39 + .../upwind_hflux_miura_stencil_01.py | 63 ++ .../advection/upwind_vflux_ppm_stencil_01.py | 40 + .../advection/v_limit_prbl_sm_stencil_01.py | 47 ++ .../advection/v_limit_prbl_sm_stencil_02.py | 60 ++ .../advection/vert_adv_stencil_01.py | 58 ++ .../atmosphere/dycore/compute_airmass.py | 36 + .../dycore/tests/test_compute_airmass.py | 40 + .../src/icon4py/model/common/dimension.py | 7 + .../model/common/test_utils/simple_mesh.py | 31 + model/requirements-dev.txt | 1 + model/requirements.txt | 1 + model/tox.ini | 4 +- pytest.ini | 2 + requirements-dev.txt | 1 + requirements.txt | 1 + tools/pyproject.toml | 1 + tools/requirements-dev.txt | 1 + tools/requirements.txt | 1 + tox.ini | 4 +- 113 files changed, 10553 insertions(+), 8 deletions(-) create mode 100644 model/atmosphere/advection/.flake8 create mode 100644 model/atmosphere/advection/.pre-commit-config.yaml create mode 100644 model/atmosphere/advection/README.md create mode 100644 model/atmosphere/advection/advection_tests/__init__.py create mode 100644 model/atmosphere/advection/advection_tests/conftest.py create mode 100644 model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_02.py create mode 100644 model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_03.py create mode 100644 model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_02.py create mode 100644 model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02.py create mode 100644 model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02a.py create mode 100644 model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02b.py create mode 100644 model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02c.py create mode 100644 model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_05.py create mode 100644 model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_01a.py create mode 100644 model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_02.py create mode 100644 model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01a.py create mode 100644 model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01b.py create mode 100644 model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_02.py create mode 100644 model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_03.py create mode 100644 model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_04.py create mode 100644 model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_02.py create mode 100644 model/atmosphere/advection/advection_tests/test_hor_adv_stencil_01.py rename model/atmosphere/{dycore/tests => advection/advection_tests}/test_mo_advection_traj_btraj_compute_o1_dsl.py (94%) create mode 100644 model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py create mode 100644 model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py create mode 100644 model/atmosphere/advection/advection_tests/test_rbf_intp_edge_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py create mode 100644 model/atmosphere/advection/advection_tests/test_recon_lsq_cell_l_svd_stencil.py create mode 100644 model/atmosphere/advection/advection_tests/test_set_zero_c.py create mode 100644 model/atmosphere/advection/advection_tests/test_set_zero_c_k.py create mode 100644 model/atmosphere/advection/advection_tests/test_step_advection_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_step_advection_stencil_02.py create mode 100644 model/atmosphere/advection/advection_tests/test_step_advection_stencil_03.py create mode 100644 model/atmosphere/advection/advection_tests/test_step_advection_stencil_04.py create mode 100644 model/atmosphere/advection/advection_tests/test_upwind_hflux_miura3_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_02.py create mode 100644 model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03a.py create mode 100644 model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03b.py create mode 100644 model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_upwind_vflux_ppm_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_vert_adv_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_01.py create mode 100644 model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_02.py create mode 100644 model/atmosphere/advection/pyproject.toml create mode 100644 model/atmosphere/advection/requirements-dev.txt create mode 100644 model/atmosphere/advection/requirements.txt create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__init__.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py rename model/atmosphere/{dycore/src/icon4py/model/atmosphere/dycore => advection/src/icon4py/model/atmosphere/advection}/mo_advection_traj_btraj_compute_o1_dsl.py (89%) create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/py.typed create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py create mode 100644 model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py create mode 100644 model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py create mode 100644 model/atmosphere/dycore/tests/test_compute_airmass.py diff --git a/.github/workflows/icon4py-qa.yml b/.github/workflows/icon4py-qa.yml index 55b260cad..4153ef528 100644 --- a/.github/workflows/icon4py-qa.yml +++ b/.github/workflows/icon4py-qa.yml @@ -58,4 +58,6 @@ jobs: - name: Run checks icon4py-model-atmosphere-diffusion run: | pre-commit run --config model/atmosphere/diffusion/.pre-commit-config.yaml --all-files - + - name: Run checks icon4py-model-atmosphere-advection + run: | + pre-commit run --config model/atmosphere/advection/.pre-commit-config.yaml --all-files diff --git a/model/README.md b/model/README.md index 0f076a932..9654b9fad 100644 --- a/model/README.md +++ b/model/README.md @@ -6,6 +6,7 @@ It includes the following packages: - `atmosphere/dycore`: Contains implementations of the dynamical core of the ICON model - `atmosphere/diffusion`: Contains the implementation of diffusion in the ICON model +- `atmosphere/advection`: Contains implementations of the advection component of the ICON model - `common`: Contains shared functionality that is required by multiple components. - `driver`: Contains the driving code for the model @@ -18,6 +19,7 @@ In the following example it is assumed that you have already created and activat ```bash # changing into the corresponding directory cd model/atmosphere/dycore +cd model/atmosphere/advection # installing a development version pip install -r requirements-dev.txt diff --git a/model/atmosphere/advection/.flake8 b/model/atmosphere/advection/.flake8 new file mode 100644 index 000000000..31cecff5a --- /dev/null +++ b/model/atmosphere/advection/.flake8 @@ -0,0 +1,42 @@ +[flake8] +# Some sane defaults for the code style checker flake8 +max-line-length = 100 +max-complexity = 15 +doctests = true +extend-ignore = + # Do not perform function calls in argument defaults + B008, + # Public code object needs docstring + D1, + # Disable dargling errors by default + DAR, + # Whitespace before ':' (black formatter breaks this sometimes) + E203, + # Line too long (using Bugbear's B950 warning) + E501, + # Line break occurred before a binary operator + W503 + +exclude = + .eggs, + .gt_cache, + .ipynb_checkpoints, + .tox, + _local_, + build, + dist, + docs, + _external_src, + tests/_disabled, + setup.py + +rst-roles = + py:mod, mod, + py:func, func, + py:data, data, + py:const, const, + py:class, class, + py:meth, meth, + py:attr, attr, + py:exc, exc, + py:obj, obj, diff --git a/model/atmosphere/advection/.pre-commit-config.yaml b/model/atmosphere/advection/.pre-commit-config.yaml new file mode 100644 index 000000000..a5f749584 --- /dev/null +++ b/model/atmosphere/advection/.pre-commit-config.yaml @@ -0,0 +1,114 @@ +# NOTE: pre-commit runs all hooks from the root folder of the repository, +# as regular git hooks do. Therefore, paths passed as arguments to the plugins +# should always be relative to the root folder. + +default_stages: [commit, push] +default_language_version: + python: python3.10 +minimum_pre_commit_version: 2.20.0 +files: "model/atmosphere/advection/.*" + +repos: +- repo: meta + hooks: + - id: check-hooks-apply + stages: [manual] + - id: check-useless-excludes + stages: [manual] + +- repo: https://github.com/asottile/setup-cfg-fmt + rev: v1.20.1 + hooks: + # Run only manually because it deletes comments + - id: setup-cfg-fmt + name: format setup.cfg + stages: [manual] + +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.3.0 + hooks: + - id: check-case-conflict + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: check-symlinks + - id: check-yaml + - id: debug-statements + - id: destroyed-symlinks + +- repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks + rev: v2.6.0 + hooks: + - id: pretty-format-ini + args: [--autofix] + - id: pretty-format-toml + args: [--autofix] + - id: pretty-format-yaml + args: [--autofix, --preserve-quotes, --indent, "2"] + +- repo: https://github.com/pre-commit/mirrors-prettier + rev: v3.0.0-alpha.4 + hooks: + - id: prettier + types_or: [markdown, json] + +- repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.3.0 + hooks: + - id: insert-license + name: add license for all ICON4Py Python source files + types: [python] + args: [--comment-style, "|#|", --license-filepath, model/.license_header.txt, --fuzzy-match-generates-todo] + +- repo: https://github.com/asottile/yesqa + rev: v1.3.0 + hooks: + - id: yesqa + +- repo: https://github.com/psf/black + rev: '22.3.0' + hooks: + - id: black + name: black Python formatter + args: [--config, model/atmosphere/advection/pyproject.toml] + +- repo: https://github.com/asottile/blacken-docs + rev: v1.12.1 + hooks: + - id: blacken-docs + name: black Python formatter for docstrings + additional_dependencies: [black==22.3.0] + +- repo: https://github.com/PyCQA/isort + rev: '5.12.0' + hooks: + - id: isort + args: [--config-root, model/atmosphere/advection/, --resolve-all-configs] + +- repo: https://github.com/PyCQA/flake8 + rev: '4.0.1' + hooks: + - id: flake8 + name: flake8 code style checks + additional_dependencies: + - darglint + - flake8-bugbear + - flake8-builtins + - flake8-debugger + - flake8-docstrings + - flake8-eradicate + - flake8-mutable + - pygments + args: [--config=model/atmosphere/advection/.flake8, model/atmosphere/advection/src/icon4py/] + +- repo: local + hooks: + - id: mypy + name: mypy static type checker + entry: bash -c 'echo mypy temporarily disabled' + #entry: bash -c 'cd model/atmosphere/dycore; mypy src/' -- + language: system + types_or: [python, pyi] + always_run: true + #pass_filenames: false + require_serial: true + stages: [commit] diff --git a/model/atmosphere/advection/README.md b/model/atmosphere/advection/README.md new file mode 100644 index 000000000..524a73d07 --- /dev/null +++ b/model/atmosphere/advection/README.md @@ -0,0 +1,9 @@ +# icon4py-atmosphere-advection + +## Description + +Contains code ported from ICON `src/advection`, which is the advection component of the ICON model. + +## Installation instructions + +Check the `README.md` at the root of the `model` folder for installation instructions. diff --git a/model/atmosphere/advection/advection_tests/__init__.py b/model/atmosphere/advection/advection_tests/__init__.py new file mode 100644 index 000000000..15dfdb009 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/__init__.py @@ -0,0 +1,12 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later diff --git a/model/atmosphere/advection/advection_tests/conftest.py b/model/atmosphere/advection/advection_tests/conftest.py new file mode 100644 index 000000000..7ee115147 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/conftest.py @@ -0,0 +1,33 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import pytest +from gt4py.next.program_processors.runners.roundtrip import executor + +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +BACKENDS = {"embedded": executor} +MESHES = {"simple_mesh": SimpleMesh()} + + +@pytest.fixture( + ids=MESHES.keys(), + params=MESHES.values(), +) +def mesh(request): + return request.param + + +@pytest.fixture(ids=BACKENDS.keys(), params=BACKENDS.values()) +def backend(request): + return request.param diff --git a/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_01.py b/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_01.py new file mode 100644 index 000000000..8d162d7b9 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_01.py @@ -0,0 +1,63 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_01 import btraj_dreg_stencil_01 +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def btraj_dreg_stencil_01_numpy( + lcounterclock: bool, + p_vn: np.array, + tangent_orientation: np.array, +): + tangent_orientation = np.expand_dims(tangent_orientation, axis=-1) + + tangent_orientation = np.broadcast_to(tangent_orientation, p_vn.shape) + + lvn_sys_pos_true = np.where(tangent_orientation * p_vn >= 0.0, True, False) + + mask_lcounterclock = np.broadcast_to(lcounterclock, p_vn.shape) + + lvn_sys_pos = np.where(mask_lcounterclock, lvn_sys_pos_true, False) + + return lvn_sys_pos + + +def test_btraj_dreg_stencil_01(): + mesh = SimpleMesh() + lcounterclock = True + p_vn = random_field(mesh, EdgeDim, KDim) + + tangent_orientation = random_field(mesh, EdgeDim) + + lvn_sys_pos = zero_field(mesh, EdgeDim, KDim, dtype=bool) + + ref = btraj_dreg_stencil_01_numpy( + lcounterclock, + np.asarray(p_vn), + np.asarray(tangent_orientation), + ) + + btraj_dreg_stencil_01( + lcounterclock, + p_vn, + tangent_orientation, + lvn_sys_pos, + offset_provider={}, + ) + + assert np.allclose(ref, lvn_sys_pos) diff --git a/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_02.py b/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_02.py new file mode 100644 index 000000000..9e59a148c --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_02.py @@ -0,0 +1,71 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_02 import btraj_dreg_stencil_02 +from icon4py.model.common.dimension import ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def btraj_dreg_stencil_02_numpy( + p_vn: np.array, + p_vt: np.array, + edge_cell_length: np.array, + p_dt: float, +): + lvn_pos = np.where(p_vn >= 0.0, True, False) + + traj_length = np.sqrt(p_vn**2 + p_vt**2) * p_dt + + edge_cell_length = np.expand_dims(edge_cell_length, axis=-1) + e2c_length = np.where(lvn_pos, edge_cell_length[:, 0], edge_cell_length[:, 1]) + + opt_famask_dsl = np.where( + traj_length > (1.25 * np.broadcast_to(e2c_length, p_vn.shape)), + int32(1), + int32(0), + ) + + return opt_famask_dsl + + +def test_btraj_dreg_stencil_02(): + mesh = SimpleMesh() + p_vn = random_field(mesh, EdgeDim, KDim) + p_vt = random_field(mesh, EdgeDim, KDim) + edge_cell_length = np.asarray(mesh.e2c, dtype=float) + edge_cell_length_new = as_1D_sparse_field(edge_cell_length, ECDim) + p_dt = 1.0 + opt_famask_dsl = zero_field(mesh, EdgeDim, KDim, dtype=int32) + + ref = btraj_dreg_stencil_02_numpy( + np.asarray(p_vn), np.asarray(p_vt), np.asarray(edge_cell_length), p_dt + ) + + btraj_dreg_stencil_02( + p_vn, + p_vt, + edge_cell_length_new, + p_dt, + opt_famask_dsl, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), + }, + ) + + assert np.allclose(ref, opt_famask_dsl) diff --git a/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_03.py b/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_03.py new file mode 100644 index 000000000..712d4ecbb --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_btraj_dreg_stencil_03.py @@ -0,0 +1,228 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.btraj_dreg_stencil_03 import btraj_dreg_stencil_03 +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, constant_field, random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def btraj_dreg_stencil_03_numpy( + p_vn: np.array, + p_vt: np.array, + cell_idx: np.array, + cell_blk: np.array, + edge_verts_1_x: np.array, + edge_verts_2_x: np.array, + edge_verts_1_y: np.array, + edge_verts_2_y: np.array, + pos_on_tplane_e_1_x: np.array, + pos_on_tplane_e_2_x: np.array, + pos_on_tplane_e_1_y: np.array, + pos_on_tplane_e_2_y: np.array, + primal_normal_cell_x: np.array, + primal_normal_cell_y: np.array, + dual_normal_cell_x: np.array, + dual_normal_cell_y: np.array, + lvn_sys_pos: np.array, + p_dt: float, +) -> tuple[np.array]: + lvn_pos = np.where(p_vn >= 0.0, True, False) + cell_idx = np.expand_dims(cell_idx, axis=-1) + cell_blk = np.expand_dims(cell_blk, axis=-1) + primal_normal_cell_x = np.expand_dims(primal_normal_cell_x, axis=-1) + dual_normal_cell_x = np.expand_dims(dual_normal_cell_x, axis=-1) + primal_normal_cell_y = np.expand_dims(primal_normal_cell_y, axis=-1) + dual_normal_cell_y = np.expand_dims(dual_normal_cell_y, axis=-1) + edge_verts_1_x = np.expand_dims(edge_verts_1_x, axis=-1) + edge_verts_1_y = np.expand_dims(edge_verts_1_y, axis=-1) + edge_verts_2_x = np.expand_dims(edge_verts_2_x, axis=-1) + edge_verts_2_y = np.expand_dims(edge_verts_2_y, axis=-1) + pos_on_tplane_e_1_x = np.expand_dims(pos_on_tplane_e_1_x, axis=-1) + pos_on_tplane_e_1_y = np.expand_dims(pos_on_tplane_e_1_y, axis=-1) + pos_on_tplane_e_2_x = np.expand_dims(pos_on_tplane_e_2_x, axis=-1) + pos_on_tplane_e_2_y = np.expand_dims(pos_on_tplane_e_2_y, axis=-1) + + p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, int32(0), int32(1)) + + depart_pts_1_x = np.broadcast_to(edge_verts_1_x, p_vn.shape) - p_vn * p_dt + depart_pts_1_y = np.broadcast_to(edge_verts_1_y, p_vn.shape) - p_vt * p_dt + depart_pts_2_x = np.broadcast_to(edge_verts_2_x, p_vn.shape) - p_vn * p_dt + depart_pts_2_y = np.broadcast_to(edge_verts_2_y, p_vn.shape) - p_vt * p_dt + + pos_on_tplane_e_x = np.where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) + pos_on_tplane_e_y = np.where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) + + pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x + pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y + pos_dreg_vert_c_2_x = np.where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + pos_dreg_vert_c_2_y = np.where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x + pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y + pos_dreg_vert_c_4_x = np.where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + pos_dreg_vert_c_4_y = np.where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + + pn_cell_1 = np.where(lvn_pos, primal_normal_cell_x[:, 0], primal_normal_cell_x[:, 1]) + pn_cell_2 = np.where(lvn_pos, primal_normal_cell_y[:, 0], primal_normal_cell_y[:, 1]) + dn_cell_1 = np.where(lvn_pos, dual_normal_cell_x[:, 0], dual_normal_cell_x[:, 1]) + dn_cell_2 = np.where(lvn_pos, dual_normal_cell_y[:, 0], dual_normal_cell_y[:, 1]) + + p_coords_dreg_v_1_lon_dsl = pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 + p_coords_dreg_v_2_lon_dsl = pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 + p_coords_dreg_v_3_lon_dsl = pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 + p_coords_dreg_v_4_lon_dsl = pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + p_coords_dreg_v_1_lat_dsl = pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 + p_coords_dreg_v_2_lat_dsl = pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 + p_coords_dreg_v_3_lat_dsl = pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 + p_coords_dreg_v_4_lat_dsl = pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 + + return ( + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl, + ) + + +def test_btraj_dreg_stencil_03(): + mesh = SimpleMesh() + + p_vn = random_field(mesh, EdgeDim, KDim) + p_vt = random_field(mesh, EdgeDim, KDim) + cell_idx = np.asarray(mesh.e2c, dtype=int32) + cell_idx_new = as_1D_sparse_field(cell_idx, ECDim) + cell_blk = constant_field(mesh, 1, EdgeDim, E2CDim, dtype=int32) + cell_blk_new = as_1D_sparse_field(cell_blk, ECDim) + + edge_verts_1_x = random_field(mesh, EdgeDim) + edge_verts_2_x = random_field(mesh, EdgeDim) + edge_verts_1_y = random_field(mesh, EdgeDim) + edge_verts_2_y = random_field(mesh, EdgeDim) + pos_on_tplane_e_1_x = random_field(mesh, EdgeDim) + pos_on_tplane_e_2_x = random_field(mesh, EdgeDim) + pos_on_tplane_e_1_y = random_field(mesh, EdgeDim) + pos_on_tplane_e_2_y = random_field(mesh, EdgeDim) + primal_normal_cell_x = random_field(mesh, EdgeDim, E2CDim) + primal_normal_cell_x_new = as_1D_sparse_field(primal_normal_cell_x, ECDim) + dual_normal_cell_x = random_field(mesh, EdgeDim, E2CDim) + dual_normal_cell_x_new = as_1D_sparse_field(dual_normal_cell_x, ECDim) + primal_normal_cell_y = random_field(mesh, EdgeDim, E2CDim) + primal_normal_cell_y_new = as_1D_sparse_field(primal_normal_cell_y, ECDim) + dual_normal_cell_y = random_field(mesh, EdgeDim, E2CDim) + dual_normal_cell_y_new = as_1D_sparse_field(dual_normal_cell_y, ECDim) + lvn_sys_pos = constant_field(mesh, True, EdgeDim, KDim, dtype=bool) + p_dt = 2.0 + p_cell_idx = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + p_cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + p_cell_blk = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + p_coords_dreg_v_1_lon_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_lon_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_lon_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_lon_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_1_lat_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_lat_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_lat_dsl = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_lat_dsl = random_field(mesh, EdgeDim, KDim) + + ( + p_cell_idx_ref, + p_cell_rel_idx_dsl_ref, + p_cell_blk_ref, + p_coords_dreg_v_1_lon_dsl_ref, + p_coords_dreg_v_2_lon_dsl_ref, + p_coords_dreg_v_3_lon_dsl_ref, + p_coords_dreg_v_4_lon_dsl_ref, + p_coords_dreg_v_1_lat_dsl_ref, + p_coords_dreg_v_2_lat_dsl_ref, + p_coords_dreg_v_3_lat_dsl_ref, + p_coords_dreg_v_4_lat_dsl_ref, + ) = btraj_dreg_stencil_03_numpy( + np.asarray(p_vn), + np.asarray(p_vt), + np.asarray(cell_idx), + np.asarray(cell_blk), + np.asarray(edge_verts_1_x), + np.asarray(edge_verts_2_x), + np.asarray(edge_verts_1_y), + np.asarray(edge_verts_2_y), + np.asarray(pos_on_tplane_e_1_x), + np.asarray(pos_on_tplane_e_2_x), + np.asarray(pos_on_tplane_e_1_y), + np.asarray(pos_on_tplane_e_2_y), + np.asarray(primal_normal_cell_x), + np.asarray(primal_normal_cell_y), + np.asarray(dual_normal_cell_x), + np.asarray(dual_normal_cell_y), + np.asarray(lvn_sys_pos), + p_dt, + ) + + btraj_dreg_stencil_03( + p_vn, + p_vt, + cell_idx_new, + cell_blk_new, + edge_verts_1_x, + edge_verts_2_x, + edge_verts_1_y, + edge_verts_2_y, + pos_on_tplane_e_1_x, + pos_on_tplane_e_2_x, + pos_on_tplane_e_1_y, + pos_on_tplane_e_2_y, + primal_normal_cell_x_new, + primal_normal_cell_y_new, + dual_normal_cell_x_new, + dual_normal_cell_y_new, + lvn_sys_pos, + p_dt, + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), + }, + ) + assert np.allclose(p_cell_idx, p_cell_idx_ref) + assert np.allclose(p_cell_rel_idx_dsl, p_cell_rel_idx_dsl_ref) + assert np.allclose(p_cell_blk, p_cell_blk_ref) + assert np.allclose(p_coords_dreg_v_1_lon_dsl, p_coords_dreg_v_1_lon_dsl_ref) + assert np.allclose(p_coords_dreg_v_2_lon_dsl, p_coords_dreg_v_2_lon_dsl_ref) + assert np.allclose(p_coords_dreg_v_3_lon_dsl, p_coords_dreg_v_3_lon_dsl_ref) + assert np.allclose(p_coords_dreg_v_4_lon_dsl, p_coords_dreg_v_4_lon_dsl_ref) + assert np.allclose(p_coords_dreg_v_1_lat_dsl, p_coords_dreg_v_1_lat_dsl_ref) + assert np.allclose(p_coords_dreg_v_2_lat_dsl, p_coords_dreg_v_2_lat_dsl_ref) + assert np.allclose(p_coords_dreg_v_3_lat_dsl, p_coords_dreg_v_3_lat_dsl_ref) + assert np.allclose(p_coords_dreg_v_4_lat_dsl, p_coords_dreg_v_4_lat_dsl_ref) diff --git a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py new file mode 100644 index 000000000..99e25a98c --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_01.py @@ -0,0 +1,799 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_01 import ( + divide_flux_area_list_stencil_01, +) +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import ( + as_1D_sparse_field, + random_field, + random_mask, + zero_field, +) +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +# FUNCTIONS +# Checking turn when travelling along three points, used to check whether lines inters. +def ccw( + p0_lon: np.array, + p0_lat: np.array, + p1_lon: np.array, + p1_lat: np.array, + p2_lon: np.array, + p2_lat: np.array, +) -> np.array: + + dx1 = p1_lon - p0_lon + dy1 = p1_lat - p0_lat + + dx2 = p2_lon - p0_lon + dy2 = p2_lat - p0_lat + + dx1dy2 = dx1 * dy2 + dy1dx2 = dy1 * dx2 + + lccw = np.where(dx1dy2 > dy1dx2, True, False) + ccw_out = np.where(lccw, int32(1), int32(-1)) # 1: clockwise, -1: counterclockwise + return ccw_out + + +# Checks whether two lines intersect +def lintersect( + line1_p1_lon: np.array, + line1_p1_lat: np.array, + line1_p2_lon: np.array, + line1_p2_lat: np.array, + line2_p1_lon: np.array, + line2_p1_lat: np.array, + line2_p2_lon: np.array, + line2_p2_lat: np.array, +) -> np.array: + + intersect1 = ccw( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + ) * ccw( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p2_lon, + line2_p2_lat, + ) + intersect2 = ccw( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p1_lon, + line1_p1_lat, + ) * ccw( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p2_lon, + line1_p2_lat, + ) + lintersect_out = np.where((intersect1 + intersect2) == -2, True, False) + + return lintersect_out + + +# Compute intersection point of two lines in 2D +def line_intersect( + line1_p1_lon: np.array, + line1_p1_lat: np.array, + line1_p2_lon: np.array, + line1_p2_lat: np.array, + line2_p1_lon: np.array, + line2_p1_lat: np.array, + line2_p2_lon: np.array, + line2_p2_lat: np.array, +) -> tuple[np.array]: + + m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) + m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) + + intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) + intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) + + return intersect_1, intersect_2 + + +def divide_flux_area_list_stencil_01_numpy( + e2c: np.array, + famask_int: np.array, + p_vn: np.array, + ptr_v3_lon: np.array, + ptr_v3_lat: np.array, + tangent_orientation_dsl: np.array, + dreg_patch0_1_lon_dsl: np.array, + dreg_patch0_1_lat_dsl: np.array, + dreg_patch0_2_lon_dsl: np.array, + dreg_patch0_2_lat_dsl: np.array, + dreg_patch0_3_lon_dsl: np.array, + dreg_patch0_3_lat_dsl: np.array, + dreg_patch0_4_lon_dsl: np.array, + dreg_patch0_4_lat_dsl: np.array, +): + ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) + ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) + ptr_v3_lon_e = np.expand_dims(ptr_v3_lon, axis=-1) + ptr_v3_lat_e = np.expand_dims(ptr_v3_lat, axis=-1) + tangent_orientation_dsl = np.expand_dims(tangent_orientation_dsl, axis=-1) + + arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl + arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl + arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl + arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl + depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl # indices have to be switched so that dep 1 belongs to arr 1 (and d2->a2) + depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl + depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl + depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl + + lvn_pos = np.where(p_vn >= 0.0, True, False) + + # get flux area departure-line segment + fl_line_p1_lon = depart_pts_1_lon_dsl + fl_line_p1_lat = depart_pts_1_lat_dsl + fl_line_p2_lon = depart_pts_2_lon_dsl + fl_line_p2_lat = depart_pts_2_lat_dsl + + # get triangle edge 1 (A1V3) + tri_line1_p1_lon = arrival_pts_1_lon_dsl + tri_line1_p1_lat = arrival_pts_1_lat_dsl + tri_line1_p2_lon = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + ) + tri_line1_p2_lat = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + ) + + # get triangle edge 2 (A2V3) + tri_line2_p1_lon = arrival_pts_2_lon_dsl + tri_line2_p1_lat = arrival_pts_2_lat_dsl + tri_line2_p2_lon = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lon_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lon_e[:, 1], p_vn.shape), + ) + tri_line2_p2_lat = np.where( + lvn_pos, + np.broadcast_to(ptr_v3_lat_e[:, 0], p_vn.shape), + np.broadcast_to(ptr_v3_lat_e[:, 1], p_vn.shape), + ) + + # Create first mask does departure-line segment intersects with A1V3 + lintersect_line1 = lintersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Create first mask does departure-line segment intersects with A2V3 + lintersect_line2 = lintersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + lvn_sys_pos = np.where( + (p_vn * np.broadcast_to(tangent_orientation_dsl, p_vn.shape)) >= 0.0, + True, + False, + ) + famask_bool = np.where(famask_int == int32(1), True, False) + # ------------------------------------------------- Case 1 + mask_case1 = np.logical_and.reduce([lintersect_line1, lintersect_line2, famask_bool]) + ps1_x, ps1_y = line_intersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + ps2_x, ps2_y = line_intersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + # Case 1 - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case1, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case1, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case1, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 1 - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_1_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_4_lon_vmask = np.where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_4_lat_vmask = np.where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 + ) + # Case 1 - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_1_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_4_lon_vmask = np.where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_4_lat_vmask = np.where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_2_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case1, np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 + ) + + # ------------------------------------------------- Case 2a + mask_case2a = np.logical_and.reduce( + [lintersect_line1, np.logical_not(lintersect_line2), famask_bool] + ) + # Case 2a - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 2a - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask) + dreg_patch1_2_lon_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), + dreg_patch1_2_lon_vmask, + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), + dreg_patch1_3_lon_vmask, + ) + dreg_patch1_3_lat_vmask = np.where( + mask_case2a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), + dreg_patch1_3_lat_vmask, + ) + # Case 2a - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) + + # -------------------------------------------------- Case 2b + mask_case2b = np.logical_and.reduce( + [lintersect_line2, np.logical_not(lintersect_line1), famask_bool] + ) + # Case 2b - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 2b - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case2b, 0.0, dreg_patch1_4_lat_vmask) + # Case 2b - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask) + dreg_patch2_2_lon_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), + dreg_patch2_2_lon_vmask, + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), + dreg_patch2_3_lon_vmask, + ) + dreg_patch2_3_lat_vmask = np.where( + mask_case2b, + np.where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), + dreg_patch2_3_lat_vmask, + ) + + # flux area edge 1 and 2 + fl_e1_p1_lon = arrival_pts_1_lon_dsl + fl_e1_p1_lat = arrival_pts_1_lat_dsl + fl_e1_p2_lon = depart_pts_1_lon_dsl + fl_e1_p2_lat = depart_pts_1_lat_dsl + fl_e2_p1_lon = arrival_pts_2_lon_dsl + fl_e2_p1_lat = arrival_pts_2_lat_dsl + fl_e2_p2_lon = depart_pts_2_lon_dsl + fl_e2_p2_lat = depart_pts_2_lat_dsl + + # ----------------------------------------------- Case 3a + # Check whether flux area edge 2 intersects with triangle edge 1 + lintersect_e2_line1 = lintersect( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + mask_case3a = np.logical_and(lintersect_e2_line1, famask_bool) + pi1_x, pi1_y = line_intersect( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Case 3a - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = np.where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 3a - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), + dreg_patch1_2_lon_vmask, + ) + dreg_patch1_2_lat_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = np.where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), + dreg_patch1_4_lon_vmask, + ) + dreg_patch1_4_lat_vmask = np.where( + mask_case3a, + np.where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), + dreg_patch1_4_lat_vmask, + ) + # Case 3a - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = np.where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) + + # ------------------------------------------------ Case 3b + # Check whether flux area edge 1 intersects with triangle edge 2 + lintersect_e1_line2 = lintersect( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + mask_case3b = lintersect_e1_line2 & famask_bool + pi2_x, pi2_y = line_intersect( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + # Case 3b - patch 0 + dreg_patch0_1_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = np.where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = np.where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_2_lon_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), + dreg_patch0_3_lon_dsl, + ) + dreg_patch0_3_lat_dsl = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), + dreg_patch0_3_lat_dsl, + ) + # Case 3b - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) + # Case 3b - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), + dreg_patch2_2_lon_vmask, + ) + dreg_patch2_2_lat_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = np.where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), + dreg_patch2_4_lon_vmask, + ) + dreg_patch2_4_lat_vmask = np.where( + mask_case3b, + np.where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), + dreg_patch2_4_lat_vmask, + ) + + # --------------------------------------------- Case 4 + # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions + # can't be overwritten by this new condition. + indices_previously_matched = np.logical_or.reduce( + [mask_case3b, mask_case3a, mask_case2b, mask_case2a, mask_case1] + ) + # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices + mask_case4 = np.logical_and.reduce([famask_bool, np.logical_not(indices_previously_matched)]) + # Case 4 - patch 0 - no change + # Case 4 - patch 1 + dreg_patch1_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = np.where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) + # Case 4 - patch 2 + dreg_patch2_1_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = np.where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = np.where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ) + + +@pytest.mark.slow_tests +def test_divide_flux_area_list_stencil_01(): + mesh = SimpleMesh() + + famask_int = random_mask(mesh, EdgeDim, KDim, dtype=int32) + p_vn = random_field(mesh, EdgeDim, KDim) + ptr_v3_lon = random_field(mesh, EdgeDim, E2CDim) + ptr_v3_lon_field = as_1D_sparse_field(ptr_v3_lon, ECDim) + ptr_v3_lat = random_field(mesh, EdgeDim, E2CDim) + ptr_v3_lat_field = as_1D_sparse_field(ptr_v3_lat, ECDim) + tangent_orientation_dsl = random_field(mesh, EdgeDim) + dreg_patch0_1_lon_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_1_lat_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_2_lon_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_2_lat_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_3_lon_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_3_lat_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_4_lon_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch0_4_lat_dsl = random_field(mesh, EdgeDim, KDim) + dreg_patch1_1_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_1_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_2_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_2_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_3_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_3_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_4_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch1_4_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_1_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_1_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_2_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_2_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_3_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_3_lat_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_4_lon_vmask = zero_field(mesh, EdgeDim, KDim) + dreg_patch2_4_lat_vmask = zero_field(mesh, EdgeDim, KDim) + + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ref_12, + ref_13, + ref_14, + ref_15, + ref_16, + ref_17, + ref_18, + ref_19, + ref_20, + ref_21, + ref_22, + ref_23, + ref_24, + ) = divide_flux_area_list_stencil_01_numpy( + mesh.e2c, + np.asarray(famask_int), + np.asarray(p_vn), + np.asarray(ptr_v3_lon), + np.asarray(ptr_v3_lat), + np.asarray(tangent_orientation_dsl), + np.asarray(dreg_patch0_1_lon_dsl), + np.asarray(dreg_patch0_1_lat_dsl), + np.asarray(dreg_patch0_2_lon_dsl), + np.asarray(dreg_patch0_2_lat_dsl), + np.asarray(dreg_patch0_3_lon_dsl), + np.asarray(dreg_patch0_3_lat_dsl), + np.asarray(dreg_patch0_4_lon_dsl), + np.asarray(dreg_patch0_4_lat_dsl), + ) + + divide_flux_area_list_stencil_01( + famask_int, + p_vn, + ptr_v3_lon_field, + ptr_v3_lat_field, + tangent_orientation_dsl, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), + }, + ) + assert np.allclose(dreg_patch0_1_lon_dsl, ref_1) + assert np.allclose(dreg_patch0_1_lat_dsl, ref_2) + assert np.allclose(dreg_patch0_2_lon_dsl, ref_3) + assert np.allclose(dreg_patch0_2_lat_dsl, ref_4) + assert np.allclose(dreg_patch0_3_lon_dsl, ref_5) + assert np.allclose(dreg_patch0_3_lat_dsl, ref_6) + assert np.allclose(dreg_patch0_4_lon_dsl, ref_7) + assert np.allclose(dreg_patch0_4_lat_dsl, ref_8) + assert np.allclose(dreg_patch1_1_lon_vmask, ref_9) + assert np.allclose(dreg_patch1_1_lat_vmask, ref_10) + assert np.allclose(dreg_patch1_2_lon_vmask, ref_11) + assert np.allclose(dreg_patch1_2_lat_vmask, ref_12) + assert np.allclose(dreg_patch1_3_lon_vmask, ref_13) + assert np.allclose(dreg_patch1_3_lat_vmask, ref_14) + assert np.allclose(dreg_patch1_4_lon_vmask, ref_15) + assert np.allclose(dreg_patch1_4_lat_vmask, ref_16) + assert np.allclose(dreg_patch2_1_lon_vmask, ref_17) + assert np.allclose(dreg_patch2_1_lat_vmask, ref_18) + assert np.allclose(dreg_patch2_2_lon_vmask, ref_19) + assert np.allclose(dreg_patch2_2_lat_vmask, ref_20) + assert np.allclose(dreg_patch2_3_lon_vmask, ref_21) + assert np.allclose(dreg_patch2_3_lat_vmask, ref_22) + assert np.allclose(dreg_patch2_4_lon_vmask, ref_23) + assert np.allclose(dreg_patch2_4_lat_vmask, ref_24) diff --git a/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_02.py new file mode 100644 index 000000000..7423360a2 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_divide_flux_area_list_stencil_02.py @@ -0,0 +1,333 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.divide_flux_area_list_stencil_02 import ( + divide_flux_area_list_stencil_02, +) +from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, random_mask +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def divide_flux_area_list_stencil_02_numpy( + e2c: np.array, + famask_int: np.array, + p_vn: np.array, + bf_cc_patch1_lon: np.array, + bf_cc_patch1_lat: np.array, + bf_cc_patch2_lon: np.array, + bf_cc_patch2_lat: np.array, + butterfly_idx_patch1_vnpos: np.array, + butterfly_idx_patch1_vnneg: np.array, + butterfly_blk_patch1_vnpos: np.array, + butterfly_blk_patch1_vnneg: np.array, + butterfly_idx_patch2_vnpos: np.array, + butterfly_idx_patch2_vnneg: np.array, + butterfly_blk_patch2_vnpos: np.array, + butterfly_blk_patch2_vnneg: np.array, + dreg_patch1_1_lon_vmask: np.array, + dreg_patch1_1_lat_vmask: np.array, + dreg_patch1_2_lon_vmask: np.array, + dreg_patch1_2_lat_vmask: np.array, + dreg_patch1_3_lon_vmask: np.array, + dreg_patch1_3_lat_vmask: np.array, + dreg_patch1_4_lon_vmask: np.array, + dreg_patch1_4_lat_vmask: np.array, + dreg_patch2_1_lon_vmask: np.array, + dreg_patch2_1_lat_vmask: np.array, + dreg_patch2_2_lon_vmask: np.array, + dreg_patch2_2_lat_vmask: np.array, + dreg_patch2_3_lon_vmask: np.array, + dreg_patch2_3_lat_vmask: np.array, + dreg_patch2_4_lon_vmask: np.array, + dreg_patch2_4_lat_vmask: np.array, +): + famask_bool = np.where(famask_int == int32(1), True, False) + lvn_pos = np.where(p_vn >= np.broadcast_to(0.0, p_vn.shape), True, False) + # Translation of patch 1 and patch 2 in system relative to respective cell + bf_cc_patch1_lon_e = np.expand_dims(bf_cc_patch1_lon, axis=-1) + bf_cc_patch1_lat_e = np.expand_dims(bf_cc_patch1_lat, axis=-1) + bf_cc_patch2_lon_e = np.expand_dims(bf_cc_patch2_lon, axis=-1) + bf_cc_patch2_lat_e = np.expand_dims(bf_cc_patch2_lat, axis=-1) + + bf_cc_patch1_lon = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch1_lon_e[:, 0], bf_cc_patch1_lon_e[:, 1]), + 0.0, + ) + bf_cc_patch1_lat = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch1_lat_e[:, 0], bf_cc_patch1_lat_e[:, 1]), + 0.0, + ) + bf_cc_patch2_lon = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch2_lon_e[:, 0], bf_cc_patch2_lon_e[:, 1]), + 0.0, + ) + bf_cc_patch2_lat = np.where( + famask_bool, + np.where(lvn_pos, bf_cc_patch2_lat_e[:, 0], bf_cc_patch2_lat_e[:, 1]), + 0.0, + ) + + # patch1 in translated system + dreg_patch1_1_lon_vmask = dreg_patch1_1_lon_vmask - bf_cc_patch1_lon + dreg_patch1_1_lat_vmask = dreg_patch1_1_lat_vmask - bf_cc_patch1_lat + dreg_patch1_2_lon_vmask = dreg_patch1_2_lon_vmask - bf_cc_patch1_lon + dreg_patch1_2_lat_vmask = dreg_patch1_2_lat_vmask - bf_cc_patch1_lat + dreg_patch1_3_lon_vmask = dreg_patch1_3_lon_vmask - bf_cc_patch1_lon + dreg_patch1_3_lat_vmask = dreg_patch1_3_lat_vmask - bf_cc_patch1_lat + dreg_patch1_4_lon_vmask = dreg_patch1_4_lon_vmask - bf_cc_patch1_lon + dreg_patch1_4_lat_vmask = dreg_patch1_4_lat_vmask - bf_cc_patch1_lat + # patch2 in translated system + dreg_patch2_1_lon_vmask = dreg_patch2_1_lon_vmask - bf_cc_patch2_lon + dreg_patch2_1_lat_vmask = dreg_patch2_1_lat_vmask - bf_cc_patch2_lat + dreg_patch2_2_lon_vmask = dreg_patch2_2_lon_vmask - bf_cc_patch2_lon + dreg_patch2_2_lat_vmask = dreg_patch2_2_lat_vmask - bf_cc_patch2_lat + dreg_patch2_3_lon_vmask = dreg_patch2_3_lon_vmask - bf_cc_patch2_lon + dreg_patch2_3_lat_vmask = dreg_patch2_3_lat_vmask - bf_cc_patch2_lat + dreg_patch2_4_lon_vmask = dreg_patch2_4_lon_vmask - bf_cc_patch2_lon + dreg_patch2_4_lat_vmask = dreg_patch2_4_lat_vmask - bf_cc_patch2_lat + + # Store global index of the underlying grid cell + # Adapt dimensions to fit ofr multiple levels + butterfly_idx_patch1_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch1_vnpos, axis=-1), p_vn.shape + ) + butterfly_idx_patch1_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch1_vnneg, axis=-1), p_vn.shape + ) + butterfly_idx_patch2_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch2_vnpos, axis=-1), p_vn.shape + ) + butterfly_idx_patch2_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_idx_patch2_vnneg, axis=-1), p_vn.shape + ) + butterfly_blk_patch1_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch1_vnpos, axis=-1), p_vn.shape + ) + butterfly_blk_patch1_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch1_vnneg, axis=-1), p_vn.shape + ) + butterfly_blk_patch2_vnpos_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch2_vnpos, axis=-1), p_vn.shape + ) + butterfly_blk_patch2_vnneg_3d = np.broadcast_to( + np.expand_dims(butterfly_blk_patch2_vnneg, axis=-1), p_vn.shape + ) + patch1_cell_idx_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_idx_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), + int32(0), + ) + patch1_cell_blk_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_blk_vmask = np.where( + famask_bool, + np.where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), + int32(0), + ) + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask, + patch1_cell_blk_vmask, + patch2_cell_idx_vmask, + patch2_cell_blk_vmask, + ) + + +def test_divide_flux_area_list_stencil_02(): + mesh = SimpleMesh() + + famask_int = random_mask(mesh, EdgeDim, KDim, dtype=int32) + p_vn = random_field(mesh, EdgeDim, KDim) + bf_cc_patch1_lon = random_field(mesh, EdgeDim, E2CDim) + bf_cc_patch1_lon_field = as_1D_sparse_field(bf_cc_patch1_lon, ECDim) + bf_cc_patch1_lat = random_field(mesh, EdgeDim, E2CDim) + bf_cc_patch1_lat_field = as_1D_sparse_field(bf_cc_patch1_lat, ECDim) + bf_cc_patch2_lon = random_field(mesh, EdgeDim, E2CDim) + bf_cc_patch2_lon_field = as_1D_sparse_field(bf_cc_patch2_lon, ECDim) + bf_cc_patch2_lat = random_field(mesh, EdgeDim, E2CDim) + bf_cc_patch2_lat_field = as_1D_sparse_field(bf_cc_patch2_lat, ECDim) + butterfly_idx_patch1_vnpos = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_idx_patch1_vnneg = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_blk_patch1_vnpos = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_blk_patch1_vnneg = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_idx_patch2_vnpos = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_idx_patch2_vnneg = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_blk_patch2_vnpos = random_mask(mesh, EdgeDim, dtype=int32) + butterfly_blk_patch2_vnneg = random_mask(mesh, EdgeDim, dtype=int32) + dreg_patch1_1_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_1_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_2_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_2_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_3_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_3_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_4_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch1_4_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_1_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_1_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_2_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_2_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_3_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_3_lat_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_4_lon_vmask = random_field(mesh, EdgeDim, KDim) + dreg_patch2_4_lat_vmask = random_field(mesh, EdgeDim, KDim) + patch1_cell_idx_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) + patch1_cell_blk_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) + patch2_cell_idx_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) + patch2_cell_blk_vmask = random_mask(mesh, EdgeDim, KDim, dtype=int32) + + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ref_12, + ref_13, + ref_14, + ref_15, + ref_16, + ref_17, + ref_18, + ref_19, + ref_20, + ) = divide_flux_area_list_stencil_02_numpy( + mesh.e2c, + np.asarray(famask_int), + np.asarray(p_vn), + np.asarray(bf_cc_patch1_lon), + np.asarray(bf_cc_patch1_lat), + np.asarray(bf_cc_patch2_lon), + np.asarray(bf_cc_patch2_lat), + np.asarray(butterfly_idx_patch1_vnpos), + np.asarray(butterfly_idx_patch1_vnneg), + np.asarray(butterfly_blk_patch1_vnpos), + np.asarray(butterfly_blk_patch1_vnneg), + np.asarray(butterfly_idx_patch2_vnpos), + np.asarray(butterfly_idx_patch2_vnneg), + np.asarray(butterfly_blk_patch2_vnpos), + np.asarray(butterfly_blk_patch2_vnneg), + np.asarray(dreg_patch1_1_lon_vmask), + np.asarray(dreg_patch1_1_lat_vmask), + np.asarray(dreg_patch1_2_lon_vmask), + np.asarray(dreg_patch1_2_lat_vmask), + np.asarray(dreg_patch1_3_lon_vmask), + np.asarray(dreg_patch1_3_lat_vmask), + np.asarray(dreg_patch1_4_lon_vmask), + np.asarray(dreg_patch1_4_lat_vmask), + np.asarray(dreg_patch2_1_lon_vmask), + np.asarray(dreg_patch2_1_lat_vmask), + np.asarray(dreg_patch2_2_lon_vmask), + np.asarray(dreg_patch2_2_lat_vmask), + np.asarray(dreg_patch2_3_lon_vmask), + np.asarray(dreg_patch2_3_lat_vmask), + np.asarray(dreg_patch2_4_lon_vmask), + np.asarray(dreg_patch2_4_lat_vmask), + ) + + divide_flux_area_list_stencil_02( + famask_int, + p_vn, + bf_cc_patch1_lon_field, + bf_cc_patch1_lat_field, + bf_cc_patch2_lon_field, + bf_cc_patch2_lat_field, + butterfly_idx_patch1_vnpos, + butterfly_idx_patch1_vnneg, + butterfly_blk_patch1_vnpos, + butterfly_blk_patch1_vnneg, + butterfly_idx_patch2_vnpos, + butterfly_idx_patch2_vnneg, + butterfly_blk_patch2_vnpos, + butterfly_blk_patch2_vnneg, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask, + patch1_cell_blk_vmask, + patch2_cell_idx_vmask, + patch2_cell_blk_vmask, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, mesh.n_e2c), + }, + ) + assert np.allclose(dreg_patch1_1_lon_vmask, ref_1) + assert np.allclose(dreg_patch1_1_lat_vmask, ref_2) + assert np.allclose(dreg_patch1_2_lon_vmask, ref_3) + assert np.allclose(dreg_patch1_2_lat_vmask, ref_4) + assert np.allclose(dreg_patch1_3_lon_vmask, ref_5) + assert np.allclose(dreg_patch1_3_lat_vmask, ref_6) + assert np.allclose(dreg_patch1_4_lon_vmask, ref_7) + assert np.allclose(dreg_patch1_4_lat_vmask, ref_8) + assert np.allclose(dreg_patch2_1_lon_vmask, ref_9) + assert np.allclose(dreg_patch2_1_lat_vmask, ref_10) + assert np.allclose(dreg_patch2_2_lon_vmask, ref_11) + assert np.allclose(dreg_patch2_2_lat_vmask, ref_12) + assert np.allclose(dreg_patch2_3_lon_vmask, ref_13) + assert np.allclose(dreg_patch2_3_lat_vmask, ref_14) + assert np.allclose(dreg_patch2_4_lon_vmask, ref_15) + assert np.allclose(dreg_patch2_4_lat_vmask, ref_16) + assert np.allclose(patch1_cell_idx_vmask, ref_17) + assert np.allclose(patch1_cell_blk_vmask, ref_18) + assert np.allclose(patch2_cell_idx_vmask, ref_19) + assert np.allclose(patch2_cell_blk_vmask, ref_20) diff --git a/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_01.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_01.py new file mode 100644 index 000000000..94ce83f87 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_01.py @@ -0,0 +1,93 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import _shape, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def face_val_ppm_stencil_01_numpy( + p_cc: np.array, + p_cellhgt_mc_now: np.array, + vert_idx: np.array, + elev: int32, +): + + # this is a comment: vert_idx = np.broadcast_to(vert_idx, p_cc.shape) + + # 01a + zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] + ) + zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1]) + z_slope_a = ( + p_cellhgt_mc_now[:, 1:-1] + / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 2:]) + ) * ( + (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 2:]) * zfac_m1 + ) + + # 01b + zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / ( + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2] + ) + zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]) + z_slope_b = ( + p_cellhgt_mc_now[:, 1:-1] + / (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]) + ) * ( + (2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac + + (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1 + ) + + z_slope = np.where(vert_idx[1:-1] < elev, z_slope_a, z_slope_b) + + return z_slope + + +def test_face_val_ppm_stencil_01(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + p_cellhgt_mc_now = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + vert_idx = zero_field(mesh, KDim, dtype=int32, extend={KDim: 1}) + + vert_idx = it_embedded.np_as_located_field(KDim)( + np.arange(0, _shape(mesh, KDim, extend={KDim: 1})[0], dtype=int32) + ) + elev = vert_idx[-2] + + z_slope = random_field(mesh, CellDim, KDim) + + ref = face_val_ppm_stencil_01_numpy( + np.asarray(p_cc), + np.asarray(p_cellhgt_mc_now), + np.asarray(vert_idx), + elev, + ) + + face_val_ppm_stencil_01( + p_cc, + p_cellhgt_mc_now, + vert_idx, + elev, + z_slope, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, :-1], z_slope[:, 1:-1]) diff --git a/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02.py new file mode 100644 index 000000000..f00485f9b --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02.py @@ -0,0 +1,88 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator import embedded as it_embedded + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02 import face_val_ppm_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import _shape, random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def face_val_ppm_stencil_02_numpy( + p_cc: np.array, + p_cellhgt_mc_now: np.array, + p_face_in: np.array, + vert_idx: np.array, + slev: int32, + elev: int32, + slevp1: int32, + elevp1: int32, +): + + p_face_a = p_face_in + + p_face_a[:, 1:] = p_cc[:, 1:] * (1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1])) + ( + p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:]) + ) * ((p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1]) + + p_face = np.where((vert_idx == slevp1) | (vert_idx == elev), p_face_a, p_face_in) + p_face = np.where((vert_idx == slev), p_cc, p_face) + p_face[:, 1:] = np.where((vert_idx[1:] == elevp1), p_cc[:, :-1], p_face[:, 1:]) + + return p_face + + +def test_face_val_ppm_stencil_02(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_cellhgt_mc_now = random_field(mesh, CellDim, KDim) + p_face_in = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim) + + vert_idx = it_embedded.np_as_located_field(KDim)( + np.arange(0, _shape(mesh, KDim)[0], dtype=int32) + ) + + slev = int32(1) + slevp1 = slev + int32(1) + elev = vert_idx[-3] + elevp1 = elev + int32(1) + + ref = face_val_ppm_stencil_02_numpy( + np.asarray(p_cc), + np.asarray(p_cellhgt_mc_now), + np.asarray(p_face_in), + np.asarray(vert_idx), + slev, + elev, + slevp1, + elevp1, + ) + + face_val_ppm_stencil_02( + p_cc, + p_cellhgt_mc_now, + p_face_in, + vert_idx, + slev, + elev, + slevp1, + elevp1, + p_face, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref, p_face) diff --git a/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02a.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02a.py new file mode 100644 index 000000000..46364ca6c --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02a.py @@ -0,0 +1,54 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def face_val_ppm_stencil_02a_numpy( + p_cc: np.array, + p_cellhgt_mc_now: np.array, +): + + p_face = p_cc.copy() + + p_face[:, 1:] = p_cc[:, 1:] * (1.0 - (p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1])) + ( + p_cellhgt_mc_now[:, 1:] / (p_cellhgt_mc_now[:, :-1] + p_cellhgt_mc_now[:, 1:]) + ) * ((p_cellhgt_mc_now[:, 1:] / p_cellhgt_mc_now[:, :-1]) * p_cc[:, 1:] + p_cc[:, :-1]) + + return p_face + + +def test_face_val_ppm_stencil_02a(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_cellhgt_mc_now = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim) + + ref = face_val_ppm_stencil_02a_numpy( + np.asarray(p_cc), + np.asarray(p_cellhgt_mc_now), + ) + + face_val_ppm_stencil_02a( + p_cc, + p_cellhgt_mc_now, + p_face, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, 1:], p_face[:, 1:]) diff --git a/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02b.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02b.py new file mode 100644 index 000000000..00a073f78 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02b.py @@ -0,0 +1,46 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02b import face_val_ppm_stencil_02b +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def face_val_ppm_stencil_02b_numpy( + p_cc: np.array, +): + + p_face = p_cc.copy() + + return p_face + + +def test_face_val_ppm_stencil_02b(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim) + + ref = face_val_ppm_stencil_02b_numpy( + np.asarray(p_cc), + ) + + face_val_ppm_stencil_02b( + p_cc, + p_face, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref, p_face) diff --git a/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02c.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02c.py new file mode 100644 index 000000000..12191447d --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_02c.py @@ -0,0 +1,48 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02c import face_val_ppm_stencil_02c +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def face_val_ppm_stencil_02c_numpy( + p_cc: np.array, +): + + p_face = p_cc.copy() + + p_face[:, 1:] = p_cc[:, :-1] + + return p_face + + +def test_face_val_ppm_stencil_02c(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim) + + ref = face_val_ppm_stencil_02c_numpy( + np.asarray(p_cc), + ) + + face_val_ppm_stencil_02c( + p_cc, + p_face, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, 1:], p_face[:, 1:]) diff --git a/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_05.py b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_05.py new file mode 100644 index 000000000..cd5744078 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_face_val_ppm_stencil_05.py @@ -0,0 +1,86 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.face_val_ppm_stencil_05 import face_val_ppm_stencil_05 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def face_val_ppm_stencil_05_numpy( + p_cc: np.array, + p_cellhgt_mc_now: np.array, + z_slope: np.array, +): + p_cellhgt_mc_now_k_minus_1 = p_cellhgt_mc_now[:, 1:-2] + p_cellhgt_mc_now_k_minus_2 = p_cellhgt_mc_now[:, 0:-3] + p_cellhgt_mc_now_k_plus_1 = p_cellhgt_mc_now[:, 3:] + p_cellhgt_mc_now = p_cellhgt_mc_now[:, 2:-1] + + p_cc_k_minus_1 = p_cc[:, 1:-1] + p_cc = p_cc[:, 2:] + z_slope_k_minus_1 = z_slope[:, 1:-1] + z_slope = z_slope[:, 2:] + + zgeo1 = p_cellhgt_mc_now_k_minus_1 / (p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now) + zgeo2 = 1.0 / ( + p_cellhgt_mc_now_k_minus_2 + + p_cellhgt_mc_now_k_minus_1 + + p_cellhgt_mc_now + + p_cellhgt_mc_now_k_plus_1 + ) + zgeo3 = (p_cellhgt_mc_now_k_minus_2 + p_cellhgt_mc_now_k_minus_1) / ( + 2.0 * p_cellhgt_mc_now_k_minus_1 + p_cellhgt_mc_now + ) + zgeo4 = (p_cellhgt_mc_now_k_plus_1 + p_cellhgt_mc_now) / ( + 2 * p_cellhgt_mc_now + p_cellhgt_mc_now_k_minus_1 + ) + + p_face = ( + p_cc_k_minus_1 + + zgeo1 * (p_cc - p_cc_k_minus_1) + + zgeo2 + * ( + (2 * p_cellhgt_mc_now * zgeo1) * (zgeo3 - zgeo4) * (p_cc - p_cc_k_minus_1) + - zgeo3 * p_cellhgt_mc_now_k_minus_1 * z_slope + + zgeo4 * p_cellhgt_mc_now * z_slope_k_minus_1 + ) + ) + + return p_face + + +def test_face_val_ppm_stencil_05(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_cellhgt_mc_now = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + z_slope = random_field(mesh, CellDim, KDim) + p_face = zero_field(mesh, CellDim, KDim) + + ref = face_val_ppm_stencil_05_numpy( + np.asarray(p_cc), + np.asarray(p_cellhgt_mc_now), + np.asarray(z_slope), + ) + + face_val_ppm_stencil_05( + p_cc, + p_cellhgt_mc_now, + z_slope, + p_face, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, :], p_face[:, 2:]) diff --git a/model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_01a.py new file mode 100644 index 000000000..df00eda2d --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_01a.py @@ -0,0 +1,205 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_01a import ( + hflux_ffsl_hybrid_stencil_01a, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hflux_ffsl_hybrid_stencil_01a_numpy( + e2c: np.array, + z_lsq_coeff_1: np.ndarray, + z_lsq_coeff_2: np.ndarray, + z_lsq_coeff_3: np.ndarray, + z_lsq_coeff_4: np.ndarray, + z_lsq_coeff_5: np.ndarray, + z_lsq_coeff_6: np.ndarray, + z_lsq_coeff_7: np.ndarray, + z_lsq_coeff_8: np.ndarray, + z_lsq_coeff_9: np.ndarray, + z_lsq_coeff_10: np.ndarray, + z_quad_vector_sum0_1: np.ndarray, + z_quad_vector_sum0_2: np.ndarray, + z_quad_vector_sum0_3: np.ndarray, + z_quad_vector_sum0_4: np.ndarray, + z_quad_vector_sum0_5: np.ndarray, + z_quad_vector_sum0_6: np.ndarray, + z_quad_vector_sum0_7: np.ndarray, + z_quad_vector_sum0_8: np.ndarray, + z_quad_vector_sum0_9: np.ndarray, + z_quad_vector_sum0_10: np.ndarray, + patch0_cell_rel_idx_dsl: np.ndarray, +): + + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] + z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] + z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] + z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] + z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] + z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] + z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] + + p_out_e_hybrid_1a = ( + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + * z_quad_vector_sum0_1 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + * z_quad_vector_sum0_2 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + * z_quad_vector_sum0_3 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4_e2c[:, 1], + z_lsq_coeff_4_e2c[:, 0], + ) + * z_quad_vector_sum0_4 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5_e2c[:, 1], + z_lsq_coeff_5_e2c[:, 0], + ) + * z_quad_vector_sum0_5 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6_e2c[:, 1], + z_lsq_coeff_6_e2c[:, 0], + ) + * z_quad_vector_sum0_6 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7_e2c[:, 1], + z_lsq_coeff_7_e2c[:, 0], + ) + * z_quad_vector_sum0_7 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8_e2c[:, 1], + z_lsq_coeff_8_e2c[:, 0], + ) + * z_quad_vector_sum0_8 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9_e2c[:, 1], + z_lsq_coeff_9_e2c[:, 0], + ) + * z_quad_vector_sum0_9 + + np.where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10_e2c[:, 1], + z_lsq_coeff_10_e2c[:, 0], + ) + * z_quad_vector_sum0_10 + ) + + return p_out_e_hybrid_1a + + +def test_hflux_ffsl_hybrid_stencil_01a(): + mesh = SimpleMesh() + z_lsq_coeff_1 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_2 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_3 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_4 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_5 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_6 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_7 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_8 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_9 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_10 = random_field(mesh, CellDim, KDim) + z_quad_vector_sum0_1 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_2 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_3 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_4 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_5 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_6 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_7 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_8 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_9 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum0_10 = random_field(mesh, EdgeDim, KDim) + patch0_cell_rel_idx_dsl = constant_field(mesh, 1, EdgeDim, KDim, dtype=int32) + p_out_e_hybrid_1a = zero_field(mesh, EdgeDim, KDim) + + ref = hflux_ffsl_hybrid_stencil_01a_numpy( + mesh.e2c, + np.asarray(z_lsq_coeff_1), + np.asarray(z_lsq_coeff_2), + np.asarray(z_lsq_coeff_3), + np.asarray(z_lsq_coeff_4), + np.asarray(z_lsq_coeff_5), + np.asarray(z_lsq_coeff_6), + np.asarray(z_lsq_coeff_7), + np.asarray(z_lsq_coeff_8), + np.asarray(z_lsq_coeff_9), + np.asarray(z_lsq_coeff_10), + np.asarray(z_quad_vector_sum0_1), + np.asarray(z_quad_vector_sum0_2), + np.asarray(z_quad_vector_sum0_3), + np.asarray(z_quad_vector_sum0_4), + np.asarray(z_quad_vector_sum0_5), + np.asarray(z_quad_vector_sum0_6), + np.asarray(z_quad_vector_sum0_7), + np.asarray(z_quad_vector_sum0_8), + np.asarray(z_quad_vector_sum0_9), + np.asarray(z_quad_vector_sum0_10), + np.asarray(patch0_cell_rel_idx_dsl), + ) + + hflux_ffsl_hybrid_stencil_01a( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + z_lsq_coeff_4, + z_lsq_coeff_5, + z_lsq_coeff_6, + z_lsq_coeff_7, + z_lsq_coeff_8, + z_lsq_coeff_9, + z_lsq_coeff_10, + z_quad_vector_sum0_1, + z_quad_vector_sum0_2, + z_quad_vector_sum0_3, + z_quad_vector_sum0_4, + z_quad_vector_sum0_5, + z_quad_vector_sum0_6, + z_quad_vector_sum0_7, + z_quad_vector_sum0_8, + z_quad_vector_sum0_9, + z_quad_vector_sum0_10, + patch0_cell_rel_idx_dsl, + p_out_e_hybrid_1a, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + + assert np.allclose(p_out_e_hybrid_1a, ref) diff --git a/model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_02.py new file mode 100644 index 000000000..f4241137d --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hflux_ffsl_hybrid_stencil_02.py @@ -0,0 +1,54 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.hflux_ffsl_hybrid_stencil_02 import ( + hflux_ffsl_hybrid_stencil_02, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hflux_ffsl_hybrid_stencil_02_numpy( + p_out_e_hybrid_2: np.ndarray, + p_mass_flx_e: np.ndarray, + z_dreg_area: np.ndarray, +): + + p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area + + return p_out_e_hybrid_2 + + +def test_hflux_ffsl_hybrid_stencil_02(): + mesh = SimpleMesh() + p_out_e_hybrid_2 = random_field(mesh, EdgeDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + z_dreg_area = random_field(mesh, EdgeDim, KDim) + + ref = hflux_ffsl_hybrid_stencil_02_numpy( + np.asarray(p_out_e_hybrid_2), + np.asarray(p_mass_flx_e), + np.asarray(z_dreg_area), + ) + + hflux_ffsl_hybrid_stencil_02( + p_out_e_hybrid_2, + p_mass_flx_e, + z_dreg_area, + offset_provider={}, + ) + + assert np.allclose(p_out_e_hybrid_2, ref) diff --git a/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01a.py new file mode 100644 index 000000000..74c1bf9d7 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01a.py @@ -0,0 +1,69 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01a import ( + hflx_limiter_mo_stencil_01a, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hflx_limiter_mo_stencil_01a_numpy( + e2c: np.array, + p_mflx_tracer_h: np.ndarray, + p_mass_flx_e: np.ndarray, + p_cc: np.ndarray, +): + p_cc_e2c = p_cc[e2c] + + z_mflx_low = 0.5 * ( + p_mass_flx_e * (p_cc_e2c[:, 0] + p_cc_e2c[:, 1]) + - np.absolute(p_mass_flx_e) * (p_cc_e2c[:, 1] - p_cc_e2c[:, 0]) + ) + + z_anti = p_mflx_tracer_h - z_mflx_low + + return (z_mflx_low, z_anti) + + +def test_hflx_limiter_mo_stencil_01a(): + mesh = SimpleMesh() + p_mflx_tracer_h = random_field(mesh, EdgeDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + p_cc = random_field(mesh, CellDim, KDim) + z_mflx_low = zero_field(mesh, EdgeDim, KDim) + z_anti = zero_field(mesh, EdgeDim, KDim) + + ref_1, ref_2 = hflx_limiter_mo_stencil_01a_numpy( + mesh.e2c, + np.asarray(p_mflx_tracer_h), + np.asarray(p_mass_flx_e), + np.asarray(p_cc), + ) + + hflx_limiter_mo_stencil_01a( + p_mflx_tracer_h, + p_mass_flx_e, + p_cc, + z_mflx_low, + z_anti, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + + assert np.allclose(z_mflx_low, ref_1) + assert np.allclose(z_anti, ref_2) diff --git a/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01b.py new file mode 100644 index 000000000..7906fc3cf --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_01b.py @@ -0,0 +1,122 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_01b import ( + hflx_limiter_mo_stencil_01b, +) +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hflx_limiter_mo_stencil_01b_numpy( + c2e: np.array, + geofac_div: np.ndarray, + p_rhodz_now: np.ndarray, + p_rhodz_new: np.ndarray, + z_mflx_low: np.ndarray, + z_anti: np.ndarray, + p_cc: np.ndarray, + p_dtime: float, +): + z_anti_c2e = z_anti[c2e] + geofac_div = np.expand_dims(geofac_div, axis=-1) + + zero_array = np.zeros(p_rhodz_now.shape) + + z_mflx_anti_1 = p_dtime * geofac_div[:, 0] / p_rhodz_new * z_anti_c2e[:, 0] + z_mflx_anti_2 = p_dtime * geofac_div[:, 1] / p_rhodz_new * z_anti_c2e[:, 1] + z_mflx_anti_3 = p_dtime * geofac_div[:, 2] / p_rhodz_new * z_anti_c2e[:, 2] + + z_mflx_anti_in = -1.0 * ( + np.minimum(zero_array, z_mflx_anti_1) + + np.minimum(zero_array, z_mflx_anti_2) + + np.minimum(zero_array, z_mflx_anti_3) + ) + + z_mflx_anti_out = ( + np.maximum(zero_array, z_mflx_anti_1) + + np.maximum(zero_array, z_mflx_anti_2) + + np.maximum(zero_array, z_mflx_anti_3) + ) + + z_fluxdiv_c = np.sum(z_mflx_low[c2e] * geofac_div, axis=1) + + z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new + z_tracer_max = np.maximum(p_cc, z_tracer_new_low) + z_tracer_min = np.minimum(p_cc, z_tracer_new_low) + + return ( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + ) + + +def test_hflx_limiter_mo_stencil_01b(): + mesh = SimpleMesh() + + geofac_div = random_field(mesh, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + p_rhodz_now = random_field(mesh, CellDim, KDim) + p_rhodz_new = random_field(mesh, CellDim, KDim) + z_mflx_low = random_field(mesh, EdgeDim, KDim) + z_anti = random_field(mesh, EdgeDim, KDim) + p_cc = random_field(mesh, CellDim, KDim) + p_dtime = 5.0 + z_mflx_anti_in = random_field(mesh, CellDim, KDim) + z_mflx_anti_out = random_field(mesh, CellDim, KDim) + z_tracer_new_low = random_field(mesh, CellDim, KDim) + z_tracer_max = random_field(mesh, CellDim, KDim) + z_tracer_min = random_field(mesh, CellDim, KDim) + + ref_1, ref_2, ref_3, ref_4, ref_5 = hflx_limiter_mo_stencil_01b_numpy( + mesh.c2e, + np.asarray(geofac_div), + np.asarray(p_rhodz_now), + np.asarray(p_rhodz_new), + np.asarray(z_mflx_low), + np.asarray(z_anti), + np.asarray(p_cc), + np.asarray(p_dtime), + ) + + hflx_limiter_mo_stencil_01b( + geofac_div_new, + p_rhodz_now, + p_rhodz_new, + z_mflx_low, + z_anti, + p_cc, + p_dtime, + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + offset_provider={ + "C2E": mesh.get_c2e_offset_provider(), + "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), + }, + ) + + assert np.allclose(z_mflx_anti_in, ref_1) + assert np.allclose(z_mflx_anti_out, ref_2) + assert np.allclose(z_tracer_new_low, ref_3) + assert np.allclose(z_tracer_max, ref_4) + assert np.allclose(z_tracer_min, ref_5) diff --git a/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_02.py new file mode 100644 index 000000000..a4a730509 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_02.py @@ -0,0 +1,129 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from numpy import int32 + +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_02 import hflx_limiter_mo_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hflx_limiter_mo_stencil_02_numpy( + refin_ctrl: np.ndarray, + p_cc: np.ndarray, + z_tracer_new_low: np.ndarray, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, + lo_bound: float, + hi_bound: float, +): + refin_ctrl = np.expand_dims(refin_ctrl, axis=1) + condition = np.logical_or( + np.equal(refin_ctrl, lo_bound * np.ones(refin_ctrl.shape, dtype=int32)), + np.equal(refin_ctrl, hi_bound * np.ones(refin_ctrl.shape, dtype=int32)), + ) + z_tracer_new_out = np.where( + condition, + np.minimum(1.1 * p_cc, np.maximum(0.9 * p_cc, z_tracer_new_low)), + z_tracer_new_low, + ) + z_tracer_max_out = np.where(condition, np.maximum(p_cc, z_tracer_new_out), z_tracer_max) + z_tracer_min_out = np.where(condition, np.minimum(p_cc, z_tracer_new_out), z_tracer_min) + return z_tracer_new_out, z_tracer_max_out, z_tracer_min_out + + +def test_hflx_limiter_mo_stencil_02_some_matching_condition(): + mesh = SimpleMesh() + + hi_bound = np.int32(1) + lo_bound = np.int32(5) + + refin_ctrl = constant_field(mesh, hi_bound, CellDim, dtype=int32) + + refin_ctrl[0:2] = np.int32(3) + + p_cc = random_field(mesh, CellDim, KDim) + z_tracer_new_low_in = random_field(mesh, CellDim, KDim) + z_tracer_max_in = random_field(mesh, CellDim, KDim) + z_tracer_min_in = random_field(mesh, CellDim, KDim) + + z_tracer_new_low = zero_field(mesh, CellDim, KDim) + z_tracer_max = zero_field(mesh, CellDim, KDim) + z_tracer_min = zero_field(mesh, CellDim, KDim) + + ref_new_low, ref_max, ref_min = hflx_limiter_mo_stencil_02_numpy( + np.asarray(refin_ctrl), + np.asarray(p_cc), + np.asarray(z_tracer_new_low_in), + np.asarray(z_tracer_max_in), + np.asarray(z_tracer_min_in), + lo_bound, + hi_bound, + ) + + hflx_limiter_mo_stencil_02( + refin_ctrl, + p_cc, + z_tracer_new_low_in, + z_tracer_max_in, + z_tracer_min_in, + lo_bound, + hi_bound, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + offset_provider={}, + ) + + assert np.allclose(z_tracer_new_low, ref_new_low) + assert np.allclose(z_tracer_max, ref_max) + assert np.allclose(z_tracer_min, ref_min) + + +def test_hflx_limiter_mo_stencil_02_none_matching_condition(): + mesh = SimpleMesh() + + hi_bound = np.int32(3) + lo_bound = np.int32(1) + + refin_ctrl = constant_field(mesh, 2, CellDim, dtype=int32) + + p_cc = random_field(mesh, CellDim, KDim) + + z_tracer_new_low_in = random_field(mesh, CellDim, KDim) + z_tracer_max_in = random_field(mesh, CellDim, KDim) + z_tracer_min_in = random_field(mesh, CellDim, KDim) + + z_tracer_new_low = zero_field(mesh, CellDim, KDim) + z_tracer_max = zero_field(mesh, CellDim, KDim) + z_tracer_min = zero_field(mesh, CellDim, KDim) + + hflx_limiter_mo_stencil_02( + refin_ctrl, + p_cc, + z_tracer_new_low_in, + z_tracer_max_in, + z_tracer_min_in, + lo_bound, + hi_bound, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + offset_provider={}, + ) + + assert np.allclose(z_tracer_new_low_in, z_tracer_new_low) + assert np.allclose(z_tracer_min_in, z_tracer_min) + assert np.allclose(z_tracer_max_in, z_tracer_max) diff --git a/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_03.py new file mode 100644 index 000000000..9677974bc --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_03.py @@ -0,0 +1,122 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at 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.advection.hflx_limiter_mo_stencil_03 import ( + hflx_limiter_mo_stencil_03, + hflx_limiter_mo_stencil_03_min_max, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hflx_limiter_mo_stencil_03_numpy( + c2e2c: np.ndarray, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, + beta_fct: float, + r_beta_fct: float, + z_mflx_anti_in: np.ndarray, + z_mflx_anti_out: np.ndarray, + z_tracer_new_low: np.ndarray, + dbl_eps: float, +): + z_max, z_min = hflx_limiter_mo_stencil_03_min_max_numpy( + c2e2c, z_tracer_max, z_tracer_min, beta_fct, r_beta_fct + ) + r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) + r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out * dbl_eps) + return r_p, r_m + + +def hflx_limiter_mo_stencil_03_min_max_numpy( + c2e2c: np.array, + z_tracer_max: np.ndarray, + z_tracer_min: np.ndarray, + beta_fct: float, + r_beta_fct: float, +) -> tuple[np.ndarray]: + z_max = beta_fct * np.maximum(np.max(z_tracer_max[c2e2c], axis=1), z_tracer_max) + z_min = r_beta_fct * np.minimum(np.min(z_tracer_min[c2e2c], axis=1), z_tracer_min) + return z_max, z_min + + +def test_hflx_diffusion_mo_stencil_03_min_max(): + mesh = SimpleMesh() + z_tracer_max = random_field(mesh, CellDim, KDim) + z_tracer_min = random_field(mesh, CellDim, KDim) + z_max = zero_field(mesh, CellDim, KDim) + z_min = zero_field(mesh, CellDim, KDim) + beta_fct = 0.9 + r_beta_fct = 0.3 + z_max_ref, z_min_ref = hflx_limiter_mo_stencil_03_min_max_numpy( + mesh.c2e2c, + np.asarray(z_tracer_max), + np.asarray(z_tracer_min), + beta_fct, + r_beta_fct, + ) + hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, + z_tracer_min, + beta_fct, + r_beta_fct, + z_max, + z_min, + offset_provider={"C2E2C": mesh.get_c2e2c_offset_provider()}, + ) + assert np.allclose(z_max, z_max_ref) + assert np.allclose(z_min, z_min_ref) + + +def test_hflx_diffusion_mo_stencil_03(): + mesh = SimpleMesh() + z_tracer_max = random_field(mesh, CellDim, KDim) + z_tracer_min = random_field(mesh, CellDim, KDim) + beta_fct = 0.4 + r_beta_fct = 0.6 + z_mflx_anti_in = random_field(mesh, CellDim, KDim) + z_mflx_anti_out = random_field(mesh, CellDim, KDim) + z_tracer_new_low = random_field(mesh, CellDim, KDim) + dbl_eps = 1e-5 + r_p = zero_field(mesh, CellDim, KDim) + r_m = zero_field(mesh, CellDim, KDim) + + r_p_ref, r_m_ref = hflx_limiter_mo_stencil_03_numpy( + mesh.c2e2c, + np.asarray(z_tracer_max), + np.asarray(z_tracer_min), + beta_fct, + r_beta_fct, + np.asarray(z_mflx_anti_in), + np.asarray(z_mflx_anti_out), + np.asarray(z_tracer_new_low), + dbl_eps, + ) + + hflx_limiter_mo_stencil_03( + z_tracer_max, + z_tracer_min, + beta_fct, + r_beta_fct, + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + dbl_eps, + r_p, + r_m, + offset_provider={"C2E2C": mesh.get_c2e2c_offset_provider()}, + ) + np.allclose(r_p_ref, r_p) + np.allclose(r_m_ref, r_m) diff --git a/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_04.py new file mode 100644 index 000000000..c2758faa6 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hflx_limiter_mo_stencil_04.py @@ -0,0 +1,59 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.hflx_limiter_mo_stencil_04 import hflx_limiter_mo_stencil_04 +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hflx_limiter_mo_stencil_04_numpy( + e2c: np.ndarray, + z_anti: np.ndarray, + r_m: np.ndarray, + r_p: np.ndarray, + z_mflx_low: np.ndarray, +): + r_frac = np.where( + z_anti >= 0, + np.minimum(r_m[e2c[:, 0]], r_p[e2c[:, 1]]), + np.minimum(r_m[e2c[:, 1]], r_p[e2c[:, 0]]), + ) + return z_mflx_low + np.minimum(1.0, r_frac) * z_anti + + +def test_hflx_limiter_mo_stencil_04(): + mesh = SimpleMesh() + z_anti = random_field(mesh, EdgeDim, KDim, low=-2.0, high=2.0) + r_m = random_field(mesh, CellDim, KDim) + r_p = random_field(mesh, CellDim, KDim) + z_mflx_low = random_field(mesh, EdgeDim, KDim) + p_mflx_tracer_h = zero_field(mesh, EdgeDim, KDim) + ref = hflx_limiter_mo_stencil_04_numpy( + mesh.e2c, + np.asarray(z_anti), + np.asarray(r_m), + np.asarray(r_p), + np.asarray(z_mflx_low), + ) + hflx_limiter_mo_stencil_04( + z_anti, + r_m, + r_p, + z_mflx_low, + p_mflx_tracer_h, + offset_provider={"E2C": mesh.get_e2c_offset_provider()}, + ) + assert np.allclose(p_mflx_tracer_h, ref) diff --git a/model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_01.py new file mode 100644 index 000000000..f837ac13b --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_01.py @@ -0,0 +1,76 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.hflx_limiter_pd_stencil_01 import hflx_limiter_pd_stencil_01 +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hflx_limiter_pd_stencil_01_numpy( + c2e: np.array, + geofac_div: np.array, + p_cc: np.array, + p_rhodz_now: np.array, + p_mflx_tracer_h: np.array, + p_dtime, + dbl_eps, +): + geofac_div = np.expand_dims(geofac_div, axis=-1) + p_m_0 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 0]] * geofac_div[:, 0] * p_dtime) + p_m_1 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 1]] * geofac_div[:, 1] * p_dtime) + p_m_2 = np.maximum(0.0, p_mflx_tracer_h[c2e[:, 2]] * geofac_div[:, 2] * p_dtime) + + p_m = p_m_0 + p_m_1 + p_m_2 + r_m = np.minimum(1.0, p_cc * p_rhodz_now / (p_m + dbl_eps)) + + return r_m + + +def test_hflx_limiter_pd_stencil_01(): + mesh = SimpleMesh() + geofac_div = random_field(mesh, CellDim, C2EDim) + p_cc = random_field(mesh, CellDim, KDim) + p_rhodz_now = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h = random_field(mesh, EdgeDim, KDim) + r_m = zero_field(mesh, CellDim, KDim) + p_dtime = np.float64(5) + dbl_eps = np.float64(1e-9) + + ref = hflx_limiter_pd_stencil_01_numpy( + mesh.c2e, + np.asarray(geofac_div), + np.asarray(p_cc), + np.asarray(p_rhodz_now), + np.asarray(p_mflx_tracer_h), + p_dtime, + dbl_eps, + ) + + hflx_limiter_pd_stencil_01( + as_1D_sparse_field(geofac_div, CEDim), + p_cc, + p_rhodz_now, + p_mflx_tracer_h, + r_m, + p_dtime, + dbl_eps, + offset_provider={ + "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), + "C2E": mesh.get_c2e_offset_provider(), + }, + ) + assert np.allclose(r_m, ref) diff --git a/model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_02.py new file mode 100644 index 000000000..c93588267 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hflx_limiter_pd_stencil_02.py @@ -0,0 +1,125 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at 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.advection.hflx_limiter_pd_stencil_02 import hflx_limiter_pd_stencil_02 +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import constant_field, random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hflx_limiter_pd_stencil_02_numpy( + e2c: np.array, + refin_ctrl: np.array, + r_m: np.array, + p_mflx_tracer_h_in: np.array, + bound, +): + r_m_e2c = r_m[e2c] + refin_ctrl = np.expand_dims(refin_ctrl, axis=-1) + p_mflx_tracer_h_out = np.where( + refin_ctrl != bound, + np.where( + p_mflx_tracer_h_in >= 0, + p_mflx_tracer_h_in * r_m_e2c[:, 0], + p_mflx_tracer_h_in * r_m_e2c[:, 1], + ), + p_mflx_tracer_h_in, + ) + return p_mflx_tracer_h_out + + +def test_hflx_limiter_pd_stencil_02_nowhere_matching_refin_ctl(): + mesh = SimpleMesh() + bound = np.int32(7) + refin_ctrl = constant_field(mesh, 4, EdgeDim, dtype=np.int32) + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + ref = hflx_limiter_pd_stencil_02_numpy( + mesh.e2c, + np.asarray(refin_ctrl), + np.asarray(r_m), + np.asarray(p_mflx_tracer_h_in), + bound, + ) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, ref) + + +def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl(): + mesh = SimpleMesh() + bound = np.int32(7) + refin_ctrl = constant_field(mesh, bound, EdgeDim, dtype=np.int32) + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, p_mflx_tracer_h_in) + + +def test_hflx_limiter_pd_stencil_02_partly_matching_refin_ctl(): + mesh = SimpleMesh() + bound = np.int32(4) + refin_ctrl = constant_field(mesh, 5, EdgeDim, dtype=np.int32) + refin_ctrl[2:6] = bound + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, p_mflx_tracer_h_in) + + +def test_hflx_limiter_pd_stencil_02_everywhere_matching_refin_ctl_does_not_change_inout_arg(): + mesh = SimpleMesh() + bound = np.int32(7) + refin_ctrl = constant_field(mesh, bound, EdgeDim, dtype=np.int32) + r_m = random_field(mesh, CellDim, KDim) + p_mflx_tracer_h_in = random_field(mesh, EdgeDim, KDim) + + hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h_in, + bound, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_mflx_tracer_h_in, p_mflx_tracer_h_in) diff --git a/model/atmosphere/advection/advection_tests/test_hor_adv_stencil_01.py b/model/atmosphere/advection/advection_tests/test_hor_adv_stencil_01.py new file mode 100644 index 000000000..2cd9e9a2c --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_hor_adv_stencil_01.py @@ -0,0 +1,80 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.hor_adv_stencil_01 import hor_adv_stencil_01 +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def hor_adv_stencil_01_numpy( + c2e: np.array, + p_mflx_tracer_h: np.array, + deepatmo_divh: np.array, + tracer_now: np.array, + rhodz_now: np.array, + rhodz_new: np.array, + geofac_div: np.array, + p_dtime, +) -> np.array: + geofac_div = np.expand_dims(geofac_div, axis=-1) + + tracer_new = ( + tracer_now * rhodz_now + - p_dtime * deepatmo_divh * np.sum(p_mflx_tracer_h[c2e] * geofac_div, axis=1) + ) / rhodz_new + + return tracer_new + + +def test_hor_adv_stencil_01(): + mesh = SimpleMesh() + + p_mflx_tracer_h = random_field(mesh, EdgeDim, KDim) + deepatmo_divh = random_field(mesh, KDim) + tracer_now = random_field(mesh, CellDim, KDim) + rhodz_now = random_field(mesh, CellDim, KDim) + rhodz_new = random_field(mesh, CellDim, KDim) + geofac_div = random_field(mesh, CellDim, C2EDim) + geofac_div_new = as_1D_sparse_field(geofac_div, CEDim) + tracer_new = random_field(mesh, CellDim, KDim) + p_dtime = np.float64(5.0) + + ref = hor_adv_stencil_01_numpy( + mesh.c2e, + np.asarray(p_mflx_tracer_h), + np.asarray(deepatmo_divh), + np.asarray(tracer_now), + np.asarray(rhodz_now), + np.asarray(rhodz_new), + np.asarray(geofac_div), + p_dtime, + ) + hor_adv_stencil_01( + p_mflx_tracer_h, + deepatmo_divh, + tracer_now, + rhodz_now, + rhodz_new, + geofac_div_new, + tracer_new, + p_dtime, + offset_provider={ + "C2E": mesh.get_c2e_offset_provider(), + "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), + }, + ) + assert np.allclose(tracer_new, ref) diff --git a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/advection_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py similarity index 94% rename from model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py rename to model/atmosphere/advection/advection_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py index becc4c26b..6f7cd10df 100644 --- a/model/atmosphere/dycore/tests/test_mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/advection_tests/test_mo_advection_traj_btraj_compute_o1_dsl.py @@ -15,7 +15,7 @@ import pytest from gt4py.next.ffront.fbuiltins import int32 -from icon4py.model.atmosphere.dycore.mo_advection_traj_btraj_compute_o1_dsl import ( +from icon4py.model.atmosphere.advection.mo_advection_traj_btraj_compute_o1_dsl import ( mo_advection_traj_btraj_compute_o1_dsl, ) from icon4py.model.common.dimension import E2CDim, ECDim, EdgeDim, KDim @@ -67,6 +67,7 @@ def reference( dual_normal_cell_2 = np.expand_dims(dual_normal_cell_2, axis=-1) p_cell_idx = np.where(lvn_pos, cell_idx[:, 0], cell_idx[:, 1]) + p_cell_rel_idx_dsl = np.where(lvn_pos, np.int32(0), np.int32(1)) p_cell_blk = np.where(lvn_pos, cell_blk[:, 0], cell_blk[:, 1]) z_ntdistv_bary_1 = -( @@ -94,6 +95,7 @@ def reference( return dict( p_cell_idx=p_cell_idx, + p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, p_cell_blk=p_cell_blk, p_distv_bary_1=p_distv_bary_1, p_distv_bary_2=p_distv_bary_2, @@ -120,6 +122,7 @@ def input_data(self, mesh): dual_normal_cell_2 = random_field(mesh, EdgeDim, E2CDim) dual_normal_cell_2_new = as_1D_sparse_field(dual_normal_cell_2, ECDim) p_cell_idx = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + p_cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) p_cell_blk = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) p_distv_bary_1 = random_field(mesh, EdgeDim, KDim) p_distv_bary_2 = random_field(mesh, EdgeDim, KDim) @@ -137,6 +140,7 @@ def input_data(self, mesh): primal_normal_cell_2=primal_normal_cell_2_new, dual_normal_cell_2=dual_normal_cell_2_new, p_cell_idx=p_cell_idx, + p_cell_rel_idx_dsl=p_cell_rel_idx_dsl, p_cell_blk=p_cell_blk, p_distv_bary_1=p_distv_bary_1, p_distv_bary_2=p_distv_bary_2, diff --git a/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py new file mode 100644 index 000000000..202669e37 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_list_stencil.py @@ -0,0 +1,505 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_list_stencil import ( + prep_gauss_quadrature_c_list_stencil, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def prep_gauss_quadrature_c_list_stencil_numpy( + famask_int: np.array, + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_dreg_area_in: np.array, +) -> tuple[np.ndarray]: + + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + z_eta_1_1 = 1.0 - eta_1 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 + z_eta_4_1 = 1.0 - eta_4 + z_eta_1_2 = 1.0 + eta_1 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 + z_eta_4_2 = 1.0 + eta_4 + z_eta_1_3 = 1.0 - zeta_1 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 + z_eta_4_3 = 1.0 - zeta_4 + z_eta_1_4 = 1.0 + zeta_1 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 + z_eta_4_4 = 1.0 + zeta_4 + + famask_bool = np.where(famask_int == int32(1), True, False) + p_coords_dreg_v_1_x = np.where(famask_bool, p_coords_dreg_v_1_x, 0.0) + p_coords_dreg_v_2_x = np.where(famask_bool, p_coords_dreg_v_2_x, 0.0) + p_coords_dreg_v_3_x = np.where(famask_bool, p_coords_dreg_v_3_x, 0.0) + p_coords_dreg_v_4_x = np.where(famask_bool, p_coords_dreg_v_4_x, 0.0) + p_coords_dreg_v_1_y = np.where(famask_bool, p_coords_dreg_v_1_y, 0.0) + p_coords_dreg_v_2_y = np.where(famask_bool, p_coords_dreg_v_2_y, 0.0) + p_coords_dreg_v_3_y = np.where(famask_bool, p_coords_dreg_v_3_y, 0.0) + p_coords_dreg_v_4_y = np.where(famask_bool, p_coords_dreg_v_4_y, 0.0) + + wgt_t_detjac_1 = np.where( + famask_bool, + dbl_eps + + z_wgt_1 + * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_2 = np.where( + famask_bool, + dbl_eps + + z_wgt_2 + * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_3 = np.where( + famask_bool, + dbl_eps + + z_wgt_3 + * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_4 = np.where( + famask_bool, + dbl_eps + + z_wgt_4 + * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + + p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area, + ) + + +@pytest.mark.slow_tests +def test_prep_gauss_quadrature_c_list_stencil(): + mesh = SimpleMesh() + + famask_int = constant_field(mesh, 1, EdgeDim, KDim, dtype=int32) + p_coords_dreg_v_1_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_1_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_y = random_field(mesh, EdgeDim, KDim) + shape_func_1_1 = 0.001 + shape_func_2_1 = 0.001 + shape_func_3_1 = 0.001 + shape_func_4_1 = 0.001 + shape_func_1_2 = 0.001 + shape_func_2_2 = 0.001 + shape_func_3_2 = 0.001 + shape_func_4_2 = 0.001 + shape_func_1_3 = 0.001 + shape_func_2_3 = 0.001 + shape_func_3_3 = 0.001 + shape_func_4_3 = 0.001 + shape_func_1_4 = 0.001 + shape_func_2_4 = 0.001 + shape_func_3_4 = 0.001 + shape_func_4_4 = 0.001 + zeta_1 = 0.002 + zeta_2 = 0.002 + zeta_3 = 0.002 + zeta_4 = 0.002 + eta_1 = 0.5 + eta_2 = 0.5 + eta_3 = 0.5 + eta_4 = 0.5 + wgt_zeta_1 = 0.003 + wgt_zeta_2 = 0.003 + wgt_eta_1 = 0.002 + wgt_eta_2 = 0.007 + dbl_eps = np.float64(0.1) + eps = 0.1 + p_dreg_area_in = random_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_1 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_2 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_3 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_4 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_5 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_6 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_7 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_8 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_9 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_10 = zero_field(mesh, EdgeDim, KDim) + p_dreg_area = zero_field(mesh, EdgeDim, KDim) + + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ) = prep_gauss_quadrature_c_list_stencil_numpy( + np.asarray(famask_int), + np.asarray(p_coords_dreg_v_1_x), + np.asarray(p_coords_dreg_v_2_x), + np.asarray(p_coords_dreg_v_3_x), + np.asarray(p_coords_dreg_v_4_x), + np.asarray(p_coords_dreg_v_1_y), + np.asarray(p_coords_dreg_v_2_y), + np.asarray(p_coords_dreg_v_3_y), + np.asarray(p_coords_dreg_v_4_y), + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + np.asarray(p_dreg_area_in), + ) + + prep_gauss_quadrature_c_list_stencil( + famask_int, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + p_dreg_area_in, + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area, + offset_provider={}, + ) + co1 = np.asarray(p_quad_vector_sum_1) + co2 = np.asarray(p_quad_vector_sum_2) + co3 = np.asarray(p_quad_vector_sum_3) + co4 = np.asarray(p_quad_vector_sum_4) + co5 = np.asarray(p_quad_vector_sum_5) + co6 = np.asarray(p_quad_vector_sum_6) + co7 = np.asarray(p_quad_vector_sum_7) + co8 = np.asarray(p_quad_vector_sum_8) + co9 = np.asarray(p_quad_vector_sum_9) + co10 = np.asarray(p_quad_vector_sum_10) + co11 = np.asarray(p_dreg_area) + assert np.allclose(ref_1, co1) + assert np.allclose(ref_2, co2) + assert np.allclose(ref_3, co3) + assert np.allclose(ref_4, co4) + assert np.allclose(ref_5, co5) + assert np.allclose(ref_6, co6) + assert np.allclose(ref_7, co7) + assert np.allclose(ref_8, co8) + assert np.allclose(ref_9, co9) + assert np.allclose(ref_10, co10) + assert np.allclose(ref_11, co11) diff --git a/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py new file mode 100644 index 000000000..876869dfb --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_prep_gauss_quadrature_c_stencil.py @@ -0,0 +1,467 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest + +from icon4py.model.atmosphere.advection.prep_gauss_quadrature_c_stencil import ( + prep_gauss_quadrature_c_stencil, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def prep_gauss_quadrature_c_stencil_numpy( + p_coords_dreg_v_1_x: np.array, + p_coords_dreg_v_2_x: np.array, + p_coords_dreg_v_3_x: np.array, + p_coords_dreg_v_4_x: np.array, + p_coords_dreg_v_1_y: np.array, + p_coords_dreg_v_2_y: np.array, + p_coords_dreg_v_3_y: np.array, + p_coords_dreg_v_4_y: np.array, + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, +) -> tuple[np.ndarray]: + + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + z_eta_1_1 = 1.0 - eta_1 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 + z_eta_4_1 = 1.0 - eta_4 + z_eta_1_2 = 1.0 + eta_1 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 + z_eta_4_2 = 1.0 + eta_4 + z_eta_1_3 = 1.0 - zeta_1 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 + z_eta_4_3 = 1.0 - zeta_4 + z_eta_1_4 = 1.0 + zeta_1 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 + z_eta_4_4 = 1.0 + zeta_4 + + wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + + z_area = p_quad_vector_sum_1 + p_dreg_area_out = np.where( + z_area >= 0.0, + np.maximum(eps, np.absolute(z_area)), + -np.maximum(eps, np.absolute(z_area)), + ) + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area_out, + ) + + +@pytest.mark.slow_tests +def test_prep_gauss_quadrature_c_stencil(): + mesh = SimpleMesh() + + p_coords_dreg_v_1_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_x = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_1_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_2_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_3_y = random_field(mesh, EdgeDim, KDim) + p_coords_dreg_v_4_y = random_field(mesh, EdgeDim, KDim) + shape_func_1_1 = 0.001 + shape_func_2_1 = 0.001 + shape_func_3_1 = 0.001 + shape_func_4_1 = 0.001 + shape_func_1_2 = 0.001 + shape_func_2_2 = 0.001 + shape_func_3_2 = 0.001 + shape_func_4_2 = 0.001 + shape_func_1_3 = 0.001 + shape_func_2_3 = 0.001 + shape_func_3_3 = 0.001 + shape_func_4_3 = 0.001 + shape_func_1_4 = 0.001 + shape_func_2_4 = 0.001 + shape_func_3_4 = 0.001 + shape_func_4_4 = 0.001 + zeta_1 = 0.002 + zeta_2 = 0.002 + zeta_3 = 0.002 + zeta_4 = 0.002 + eta_1 = 0.5 + eta_2 = 0.5 + eta_3 = 0.5 + eta_4 = 0.5 + wgt_zeta_1 = 0.003 + wgt_zeta_2 = 0.003 + wgt_eta_1 = 0.002 + wgt_eta_2 = 0.007 + dbl_eps = np.float64(0.1) + eps = 0.1 + p_quad_vector_sum_1 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_2 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_3 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_4 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_5 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_6 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_7 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_8 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_9 = zero_field(mesh, EdgeDim, KDim) + p_quad_vector_sum_10 = zero_field(mesh, EdgeDim, KDim) + p_dreg_area_out = zero_field(mesh, EdgeDim, KDim) + + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ref_11, + ) = prep_gauss_quadrature_c_stencil_numpy( + np.asarray(p_coords_dreg_v_1_x), + np.asarray(p_coords_dreg_v_2_x), + np.asarray(p_coords_dreg_v_3_x), + np.asarray(p_coords_dreg_v_4_x), + np.asarray(p_coords_dreg_v_1_y), + np.asarray(p_coords_dreg_v_2_y), + np.asarray(p_coords_dreg_v_3_y), + np.asarray(p_coords_dreg_v_4_y), + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + ) + + prep_gauss_quadrature_c_stencil( + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area_out, + offset_provider={}, + ) + co1 = np.asarray(p_quad_vector_sum_1) + co2 = np.asarray(p_quad_vector_sum_2) + co3 = np.asarray(p_quad_vector_sum_3) + co4 = np.asarray(p_quad_vector_sum_4) + co5 = np.asarray(p_quad_vector_sum_5) + co6 = np.asarray(p_quad_vector_sum_6) + co7 = np.asarray(p_quad_vector_sum_7) + co8 = np.asarray(p_quad_vector_sum_8) + co9 = np.asarray(p_quad_vector_sum_9) + co10 = np.asarray(p_quad_vector_sum_10) + co11 = np.asarray(p_dreg_area_out) + assert np.allclose(ref_1, co1) + assert np.allclose(ref_2, co2) + assert np.allclose(ref_3, co3) + assert np.allclose(ref_4, co4) + assert np.allclose(ref_5, co5) + assert np.allclose(ref_6, co6) + assert np.allclose(ref_7, co7) + assert np.allclose(ref_8, co8) + assert np.allclose(ref_9, co9) + assert np.allclose(ref_10, co10) + assert np.allclose(ref_11, co11) diff --git a/model/atmosphere/advection/advection_tests/test_rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/advection_tests/test_rbf_intp_edge_stencil_01.py new file mode 100644 index 000000000..a66751bf9 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_rbf_intp_edge_stencil_01.py @@ -0,0 +1,54 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.rbf_intp_edge_stencil_01 import rbf_intp_edge_stencil_01 +from icon4py.model.common.dimension import E2C2EDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def rbf_intp_edge_stencil_01_numpy( + e2c2e: np.array, + p_vn_in: np.array, + ptr_coeff: np.array, +) -> np.array: + + ptr_coeff = np.expand_dims(ptr_coeff, axis=-1) + p_vt_out = np.sum(p_vn_in[e2c2e] * ptr_coeff, axis=1) + return p_vt_out + + +def test_rbf_intp_edge_stencil_01(): + mesh = SimpleMesh() + + p_vn_in = random_field(mesh, EdgeDim, KDim) + ptr_coeff = random_field(mesh, EdgeDim, E2C2EDim) + p_vt_out = zero_field(mesh, EdgeDim, KDim) + + ref = rbf_intp_edge_stencil_01_numpy( + mesh.e2c2e, + np.asarray(p_vn_in), + np.asarray(ptr_coeff), + ) + + rbf_intp_edge_stencil_01( + p_vn_in, + ptr_coeff, + p_vt_out, + offset_provider={ + "E2C2E": mesh.get_e2c2e_offset_provider(), + }, + ) + assert np.allclose(p_vt_out, ref) diff --git a/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py new file mode 100644 index 000000000..7dda79538 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_c_svd_stencil.py @@ -0,0 +1,340 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +import pytest +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.recon_lsq_cell_c_svd_stencil import ( + recon_lsq_cell_c_svd_stencil, +) +from icon4py.model.common.dimension import C2E2C2E2CDim, CECECDim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def recon_lsq_cell_c_svd_stencil_numpy( + c2e2c2e2c: np.ndarray, + p_cc: np.ndarray, + lsq_pseudoinv_1: np.ndarray, + lsq_pseudoinv_2: np.ndarray, + lsq_pseudoinv_3: np.ndarray, + lsq_pseudoinv_4: np.ndarray, + lsq_pseudoinv_5: np.ndarray, + lsq_pseudoinv_6: np.ndarray, + lsq_pseudoinv_7: np.ndarray, + lsq_pseudoinv_8: np.ndarray, + lsq_pseudoinv_9: np.ndarray, + lsq_moments_1: np.ndarray, + lsq_moments_2: np.ndarray, + lsq_moments_3: np.ndarray, + lsq_moments_4: np.ndarray, + lsq_moments_5: np.ndarray, + lsq_moments_6: np.ndarray, + lsq_moments_7: np.ndarray, + lsq_moments_8: np.ndarray, + lsq_moments_9: np.ndarray, +) -> tuple[np.ndarray]: + + lsq_moments_1 = np.expand_dims(lsq_moments_1, axis=-1) + lsq_moments_2 = np.expand_dims(lsq_moments_2, axis=-1) + lsq_moments_3 = np.expand_dims(lsq_moments_3, axis=-1) + lsq_moments_4 = np.expand_dims(lsq_moments_4, axis=-1) + lsq_moments_5 = np.expand_dims(lsq_moments_5, axis=-1) + lsq_moments_6 = np.expand_dims(lsq_moments_6, axis=-1) + lsq_moments_7 = np.expand_dims(lsq_moments_7, axis=-1) + lsq_moments_8 = np.expand_dims(lsq_moments_8, axis=-1) + lsq_moments_9 = np.expand_dims(lsq_moments_9, axis=-1) + lsq_moments_1 = np.broadcast_to(lsq_moments_1, p_cc.shape) + lsq_moments_2 = np.broadcast_to(lsq_moments_2, p_cc.shape) + lsq_moments_3 = np.broadcast_to(lsq_moments_3, p_cc.shape) + lsq_moments_4 = np.broadcast_to(lsq_moments_4, p_cc.shape) + lsq_moments_5 = np.broadcast_to(lsq_moments_5, p_cc.shape) + lsq_moments_6 = np.broadcast_to(lsq_moments_6, p_cc.shape) + lsq_moments_7 = np.broadcast_to(lsq_moments_7, p_cc.shape) + lsq_moments_8 = np.broadcast_to(lsq_moments_8, p_cc.shape) + lsq_moments_9 = np.broadcast_to(lsq_moments_9, p_cc.shape) + lsq_pseudoinv_9 = np.expand_dims(lsq_pseudoinv_9, axis=-1) + lsq_pseudoinv_8 = np.expand_dims(lsq_pseudoinv_8, axis=-1) + lsq_pseudoinv_7 = np.expand_dims(lsq_pseudoinv_7, axis=-1) + lsq_pseudoinv_6 = np.expand_dims(lsq_pseudoinv_6, axis=-1) + lsq_pseudoinv_5 = np.expand_dims(lsq_pseudoinv_5, axis=-1) + lsq_pseudoinv_4 = np.expand_dims(lsq_pseudoinv_4, axis=-1) + lsq_pseudoinv_3 = np.expand_dims(lsq_pseudoinv_3, axis=-1) + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) + + p_coeff_10_dsl = ( + lsq_pseudoinv_9[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_9[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_9[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_9[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_9[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_9[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_9[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_9[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_9[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_9_dsl = ( + lsq_pseudoinv_8[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_8[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_8[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_8[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_8[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_8[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_8[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_8[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_8[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_8_dsl = ( + lsq_pseudoinv_7[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_7[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_7[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_7[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_7[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_7[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_7[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_7[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_7[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_7_dsl = ( + lsq_pseudoinv_6[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_6[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_6[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_6[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_6[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_6[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_6[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_6[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_6[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_6_dsl = ( + lsq_pseudoinv_5[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_5[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_5[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_5[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_5[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_5[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_5[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_5[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_5[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_5_dsl = ( + lsq_pseudoinv_4[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_4[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_4[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_4[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_4[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_4[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_4[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_4[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_4[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_4_dsl = ( + lsq_pseudoinv_3[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_3[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_3[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_3[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_3[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_3[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_3[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_3[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_3[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_3_dsl = ( + lsq_pseudoinv_2[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_2[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_2[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_2[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_2[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_2[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_2[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_2[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_2[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_2_dsl = ( + lsq_pseudoinv_1[:, 0] * (p_cc[c2e2c2e2c[:, 0]] - p_cc) + + lsq_pseudoinv_1[:, 1] * (p_cc[c2e2c2e2c[:, 1]] - p_cc) + + lsq_pseudoinv_1[:, 2] * (p_cc[c2e2c2e2c[:, 2]] - p_cc) + + lsq_pseudoinv_1[:, 3] * (p_cc[c2e2c2e2c[:, 3]] - p_cc) + + lsq_pseudoinv_1[:, 4] * (p_cc[c2e2c2e2c[:, 4]] - p_cc) + + lsq_pseudoinv_1[:, 5] * (p_cc[c2e2c2e2c[:, 5]] - p_cc) + + lsq_pseudoinv_1[:, 6] * (p_cc[c2e2c2e2c[:, 6]] - p_cc) + + lsq_pseudoinv_1[:, 7] * (p_cc[c2e2c2e2c[:, 7]] - p_cc) + + lsq_pseudoinv_1[:, 8] * (p_cc[c2e2c2e2c[:, 8]] - p_cc) + ) + + p_coeff_1_dsl = p_cc - ( + p_coeff_2_dsl * lsq_moments_1 + + p_coeff_3_dsl * lsq_moments_2 + + p_coeff_4_dsl * lsq_moments_3 + + p_coeff_5_dsl * lsq_moments_4 + + p_coeff_6_dsl * lsq_moments_5 + + p_coeff_7_dsl * lsq_moments_6 + + p_coeff_8_dsl * lsq_moments_7 + + p_coeff_9_dsl * lsq_moments_8 + + p_coeff_10_dsl * lsq_moments_9 + ) + return ( + p_coeff_1_dsl, + p_coeff_2_dsl, + p_coeff_3_dsl, + p_coeff_4_dsl, + p_coeff_5_dsl, + p_coeff_6_dsl, + p_coeff_7_dsl, + p_coeff_8_dsl, + p_coeff_9_dsl, + p_coeff_10_dsl, + ) + + +@pytest.mark.slow_tests +def test_recon_lsq_cell_c_svd_stencil(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + lsq_pseudoinv_1 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_2 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_3 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_4 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_5 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_6 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_7 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_8 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_9 = random_field(mesh, CellDim, C2E2C2E2CDim) + lsq_pseudoinv_1_field = as_1D_sparse_field(lsq_pseudoinv_1, CECECDim) + lsq_pseudoinv_2_field = as_1D_sparse_field(lsq_pseudoinv_2, CECECDim) + lsq_pseudoinv_3_field = as_1D_sparse_field(lsq_pseudoinv_3, CECECDim) + lsq_pseudoinv_4_field = as_1D_sparse_field(lsq_pseudoinv_4, CECECDim) + lsq_pseudoinv_5_field = as_1D_sparse_field(lsq_pseudoinv_5, CECECDim) + lsq_pseudoinv_6_field = as_1D_sparse_field(lsq_pseudoinv_6, CECECDim) + lsq_pseudoinv_7_field = as_1D_sparse_field(lsq_pseudoinv_7, CECECDim) + lsq_pseudoinv_8_field = as_1D_sparse_field(lsq_pseudoinv_8, CECECDim) + lsq_pseudoinv_9_field = as_1D_sparse_field(lsq_pseudoinv_9, CECECDim) + lsq_moments_1 = random_field(mesh, CellDim) + lsq_moments_2 = random_field(mesh, CellDim) + lsq_moments_3 = random_field(mesh, CellDim) + lsq_moments_4 = random_field(mesh, CellDim) + lsq_moments_5 = random_field(mesh, CellDim) + lsq_moments_6 = random_field(mesh, CellDim) + lsq_moments_7 = random_field(mesh, CellDim) + lsq_moments_8 = random_field(mesh, CellDim) + lsq_moments_9 = random_field(mesh, CellDim) + p_coeff_1_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_2_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_3_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_4_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_5_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_6_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_7_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_8_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_9_dsl = zero_field(mesh, CellDim, KDim) + p_coeff_10_dsl = zero_field(mesh, CellDim, KDim) + + ( + ref_1, + ref_2, + ref_3, + ref_4, + ref_5, + ref_6, + ref_7, + ref_8, + ref_9, + ref_10, + ) = recon_lsq_cell_c_svd_stencil_numpy( + mesh.c2e2c2e2c, + np.asarray(p_cc), + np.asarray(lsq_pseudoinv_1), + np.asarray(lsq_pseudoinv_2), + np.asarray(lsq_pseudoinv_3), + np.asarray(lsq_pseudoinv_4), + np.asarray(lsq_pseudoinv_5), + np.asarray(lsq_pseudoinv_6), + np.asarray(lsq_pseudoinv_7), + np.asarray(lsq_pseudoinv_8), + np.asarray(lsq_pseudoinv_9), + np.asarray(lsq_moments_1), + np.asarray(lsq_moments_2), + np.asarray(lsq_moments_3), + np.asarray(lsq_moments_4), + np.asarray(lsq_moments_5), + np.asarray(lsq_moments_6), + np.asarray(lsq_moments_7), + np.asarray(lsq_moments_8), + np.asarray(lsq_moments_9), + ) + + recon_lsq_cell_c_svd_stencil( + p_cc, + lsq_pseudoinv_1_field, + lsq_pseudoinv_2_field, + lsq_pseudoinv_3_field, + lsq_pseudoinv_4_field, + lsq_pseudoinv_5_field, + lsq_pseudoinv_6_field, + lsq_pseudoinv_7_field, + lsq_pseudoinv_8_field, + lsq_pseudoinv_9_field, + lsq_moments_1, + lsq_moments_2, + lsq_moments_3, + lsq_moments_4, + lsq_moments_5, + lsq_moments_6, + lsq_moments_7, + lsq_moments_8, + lsq_moments_9, + p_coeff_1_dsl, + p_coeff_2_dsl, + p_coeff_3_dsl, + p_coeff_4_dsl, + p_coeff_5_dsl, + p_coeff_6_dsl, + p_coeff_7_dsl, + p_coeff_8_dsl, + p_coeff_9_dsl, + p_coeff_10_dsl, + offset_provider={ + "C2E2C2E2C": mesh.get_c2e2c2e2c_offset_provider(), + "C2CECEC": StridedNeighborOffsetProvider(CellDim, CECECDim, mesh.n_c2e2c2e2c), + }, + ) + co1 = np.asarray(p_coeff_1_dsl) + co2 = np.asarray(p_coeff_2_dsl) + co3 = np.asarray(p_coeff_3_dsl) + co4 = np.asarray(p_coeff_4_dsl) + co5 = np.asarray(p_coeff_5_dsl) + co6 = np.asarray(p_coeff_6_dsl) + co7 = np.asarray(p_coeff_7_dsl) + co8 = np.asarray(p_coeff_8_dsl) + co9 = np.asarray(p_coeff_9_dsl) + co10 = np.asarray(p_coeff_10_dsl) + assert np.allclose(ref_1, co1) + assert np.allclose(ref_2, co2) + assert np.allclose(ref_3, co3) + assert np.allclose(ref_4, co4) + assert np.allclose(ref_5, co5) + assert np.allclose(ref_6, co6) + assert np.allclose(ref_7, co7) + assert np.allclose(ref_8, co8) + assert np.allclose(ref_9, co9) + assert np.allclose(ref_10, co10) diff --git a/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_l_svd_stencil.py new file mode 100644 index 000000000..5d72af698 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_recon_lsq_cell_l_svd_stencil.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 gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.recon_lsq_cell_l_svd_stencil import ( + recon_lsq_cell_l_svd_stencil, +) +from icon4py.model.common.dimension import C2E2CDim, CECDim, CellDim, KDim +from icon4py.model.common.test_utils.helpers import as_1D_sparse_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def recon_lsq_cell_l_svd_stencil_numpy( + c2e2c: np.ndarray, + p_cc: np.ndarray, + lsq_pseudoinv_1: np.ndarray, + lsq_pseudoinv_2: np.ndarray, +) -> tuple[np.ndarray]: + p_cc_e = np.expand_dims(p_cc, axis=1) + n_diff = p_cc[c2e2c] - p_cc_e + lsq_pseudoinv_2 = np.expand_dims(lsq_pseudoinv_2, axis=-1) + lsq_pseudoinv_1 = np.expand_dims(lsq_pseudoinv_1, axis=-1) + p_coeff_1 = p_cc + p_coeff_2 = np.sum(lsq_pseudoinv_1 * n_diff, axis=1) + p_coeff_3 = np.sum(lsq_pseudoinv_2 * n_diff, axis=1) + return p_coeff_1, p_coeff_2, p_coeff_3 + + +def test_recon_lsq_cell_l_svd_stencil(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + lsq_pseudoinv_1 = random_field(mesh, CellDim, C2E2CDim) + lsq_pseudoinv_1_field = as_1D_sparse_field(lsq_pseudoinv_1, CECDim) + + lsq_pseudoinv_2 = random_field(mesh, CellDim, C2E2CDim) + lsq_pseudoinv_2_field = as_1D_sparse_field(lsq_pseudoinv_2, CECDim) + p_coeff_1 = zero_field(mesh, CellDim, KDim) + p_coeff_2 = zero_field(mesh, CellDim, KDim) + p_coeff_3 = zero_field(mesh, CellDim, KDim) + + ref_1, ref_2, ref_3 = recon_lsq_cell_l_svd_stencil_numpy( + mesh.c2e2c, + np.asarray(p_cc), + np.asarray(lsq_pseudoinv_1), + np.asarray(lsq_pseudoinv_2), + ) + + recon_lsq_cell_l_svd_stencil( + p_cc, + lsq_pseudoinv_1_field, + lsq_pseudoinv_2_field, + p_coeff_1, + p_coeff_2, + p_coeff_3, + offset_provider={ + "C2E2C": mesh.get_c2e2c_offset_provider(), + "C2CEC": StridedNeighborOffsetProvider(CellDim, CECDim, mesh.n_c2e2c), + }, + ) + co1 = np.asarray(p_coeff_1) + co2 = np.asarray(p_coeff_2) + co3 = np.asarray(p_coeff_3) + assert np.allclose(ref_1, co1) + assert np.allclose(ref_2, co2) + assert np.allclose(ref_3, co3) diff --git a/model/atmosphere/advection/advection_tests/test_set_zero_c.py b/model/atmosphere/advection/advection_tests/test_set_zero_c.py new file mode 100644 index 000000000..f023b44e6 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_set_zero_c.py @@ -0,0 +1,27 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.set_zero_c import set_zero_c +from icon4py.model.common.dimension import CellDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def test_set_zero_cell_k(): + mesh = SimpleMesh() + field = random_field(mesh, CellDim) + + set_zero_c(field, offset_provider={}) + assert np.allclose(field, zero_field(mesh, CellDim)) diff --git a/model/atmosphere/advection/advection_tests/test_set_zero_c_k.py b/model/atmosphere/advection/advection_tests/test_set_zero_c_k.py new file mode 100644 index 000000000..e0f4afc4b --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_set_zero_c_k.py @@ -0,0 +1,27 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.set_zero_c_k import set_zero_c_k +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def test_set_zero_c_k(): + mesh = SimpleMesh() + field = random_field(mesh, CellDim, KDim) + + set_zero_c_k(field, offset_provider={}) + assert np.allclose(field, zero_field(mesh, CellDim, KDim)) diff --git a/model/atmosphere/advection/advection_tests/test_step_advection_stencil_01.py b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_01.py new file mode 100644 index 000000000..97ab7f488 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_01.py @@ -0,0 +1,62 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at 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.advection.step_advection_stencil_01 import step_advection_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def step_advection_stencil_01_numpy( + rhodz_ast: np.ndarray, + p_mflx_contra_v: np.ndarray, + deepatmo_divzl: np.ndarray, + deepatmo_divzu: np.ndarray, + pd_time: float, +) -> np.ndarray: + tmp = pd_time * ( + p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu + ) + return rhodz_ast + tmp + + +def test_step_advection_stencil_01(): + mesh = SimpleMesh() + rhodz_ast = random_field(mesh, CellDim, KDim) + p_mflx_contra = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + deepatmo_divzl = random_field(mesh, KDim) + deepatmo_divzu = random_field(mesh, KDim) + result = zero_field(mesh, CellDim, KDim) + p_dtime = 0.1 + + ref = step_advection_stencil_01_numpy( + np.asarray(rhodz_ast), + np.asarray(p_mflx_contra), + np.asarray(deepatmo_divzl), + np.asarray(deepatmo_divzu), + p_dtime, + ) + + step_advection_stencil_01( + rhodz_ast, + p_mflx_contra, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + result, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, :-1], result[:, :-1]) diff --git a/model/atmosphere/advection/advection_tests/test_step_advection_stencil_02.py b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_02.py new file mode 100644 index 000000000..c6604dde9 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_02.py @@ -0,0 +1,61 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.step_advection_stencil_02 import step_advection_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def step_advection_stencil_02_numpy( + rhodz_new: np.ndarray, + p_mflx_contra_v: np.ndarray, + deepatmo_divzl: np.ndarray, + deepatmo_divzu: np.ndarray, + pd_time: float, +) -> np.ndarray: + + tmp = p_mflx_contra_v[:, 1:] * deepatmo_divzl - p_mflx_contra_v[:, :-1] * deepatmo_divzu + return np.maximum(0.1 * rhodz_new, rhodz_new) - pd_time * tmp + + +def test_step_advection_stencil_02(): + mesh = SimpleMesh() + rhodz_ast = random_field(mesh, CellDim, KDim) + p_mflx_contra = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + deepatmo_divzl = random_field(mesh, KDim) + deepatmo_divzu = random_field(mesh, KDim) + result = zero_field(mesh, CellDim, KDim) + p_dtime = 0.1 + + ref = step_advection_stencil_02_numpy( + np.asarray(rhodz_ast), + np.asarray(p_mflx_contra), + np.asarray(deepatmo_divzl), + np.asarray(deepatmo_divzu), + p_dtime, + ) + + step_advection_stencil_02( + rhodz_ast, + p_mflx_contra, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + result, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(ref[:, :-1], result[:, :-1]) diff --git a/model/atmosphere/advection/advection_tests/test_step_advection_stencil_03.py b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_03.py new file mode 100644 index 000000000..87082a32d --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_03.py @@ -0,0 +1,55 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.step_advection_stencil_03 import step_advection_stencil_03 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def step_advection_stencil_03_numpy( + p_tracer_now: np.array, + p_grf_tend_tracer: np.array, + p_dtime, +) -> np.array: + + p_tracer_new = p_tracer_now + p_dtime * p_grf_tend_tracer + p_tracer_new = np.where(p_tracer_new < 0.0, 0.0, p_tracer_new) + + return p_tracer_new + + +def test_step_advection_stencil_03(): + mesh = SimpleMesh() + + p_tracer_now = random_field(mesh, CellDim, KDim) + p_grf_tend_tracer = random_field(mesh, CellDim, KDim) + p_tracer_new = random_field(mesh, CellDim, KDim) + + p_dtime = np.float64(5.0) + + ref = step_advection_stencil_03_numpy( + np.asarray(p_tracer_now), + np.asarray(p_grf_tend_tracer), + p_dtime, + ) + step_advection_stencil_03( + p_tracer_now, + p_grf_tend_tracer, + p_tracer_new, + p_dtime, + offset_provider={}, + ) + assert np.allclose(p_tracer_new, ref) diff --git a/model/atmosphere/advection/advection_tests/test_step_advection_stencil_04.py b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_04.py new file mode 100644 index 000000000..dca7722d0 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_step_advection_stencil_04.py @@ -0,0 +1,51 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.step_advection_stencil_04 import step_advection_stencil_04 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def step_advection_stencil_04_numpy( + p_tracer_now: np.array, + p_tracer_new: np.array, + p_dtime, +) -> np.array: + opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime + return opt_ddt_tracer_adv + + +def test_step_advection_stencil_04(): + mesh = SimpleMesh() + + p_tracer_now = random_field(mesh, CellDim, KDim) + p_tracer_new = random_field(mesh, CellDim, KDim) + opt_ddt_tracer_adv = zero_field(mesh, CellDim, KDim) + p_dtime = np.float64(5.0) + + ref = step_advection_stencil_04_numpy( + np.asarray(p_tracer_now), + np.asarray(p_tracer_new), + p_dtime, + ) + step_advection_stencil_04( + p_tracer_now, + p_tracer_new, + opt_ddt_tracer_adv, + p_dtime, + offset_provider={}, + ) + assert np.allclose(opt_ddt_tracer_adv, ref) diff --git a/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura3_stencil_01.py new file mode 100644 index 000000000..79a0adc68 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura3_stencil_01.py @@ -0,0 +1,217 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.upwind_hflux_miura3_stencil_01 import ( + upwind_hflux_miura3_stencil_01, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura3_stencil_01_numpy( + e2c: np.array, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + z_lsq_coeff_4: np.array, + z_lsq_coeff_5: np.array, + z_lsq_coeff_6: np.array, + z_lsq_coeff_7: np.array, + z_lsq_coeff_8: np.array, + z_lsq_coeff_9: np.array, + z_lsq_coeff_10: np.array, + z_quad_vector_sum_1: np.array, + z_quad_vector_sum_2: np.array, + z_quad_vector_sum_3: np.array, + z_quad_vector_sum_4: np.array, + z_quad_vector_sum_5: np.array, + z_quad_vector_sum_6: np.array, + z_quad_vector_sum_7: np.array, + z_quad_vector_sum_8: np.array, + z_quad_vector_sum_9: np.array, + z_quad_vector_sum_10: np.array, + z_dreg_area: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, +) -> np.array: + + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + z_lsq_coeff_4_e2c = z_lsq_coeff_4[e2c] + z_lsq_coeff_5_e2c = z_lsq_coeff_5[e2c] + z_lsq_coeff_6_e2c = z_lsq_coeff_6[e2c] + z_lsq_coeff_7_e2c = z_lsq_coeff_7[e2c] + z_lsq_coeff_8_e2c = z_lsq_coeff_8[e2c] + z_lsq_coeff_9_e2c = z_lsq_coeff_9[e2c] + z_lsq_coeff_10_e2c = z_lsq_coeff_10[e2c] + + p_out_e_miura3 = ( + ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + * z_quad_vector_sum_1 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + * z_quad_vector_sum_2 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + * z_quad_vector_sum_3 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4_e2c[:, 1], + z_lsq_coeff_4_e2c[:, 0], + ) + * z_quad_vector_sum_4 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5_e2c[:, 1], + z_lsq_coeff_5_e2c[:, 0], + ) + * z_quad_vector_sum_5 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6_e2c[:, 1], + z_lsq_coeff_6_e2c[:, 0], + ) + * z_quad_vector_sum_6 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7_e2c[:, 1], + z_lsq_coeff_7_e2c[:, 0], + ) + * z_quad_vector_sum_7 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8_e2c[:, 1], + z_lsq_coeff_8_e2c[:, 0], + ) + * z_quad_vector_sum_8 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9_e2c[:, 1], + z_lsq_coeff_9_e2c[:, 0], + ) + * z_quad_vector_sum_9 + + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10_e2c[:, 1], + z_lsq_coeff_10_e2c[:, 0], + ) + * z_quad_vector_sum_10 + ) + / z_dreg_area + * p_mass_flx_e + ) + + return p_out_e_miura3 + + +def test_upwind_hflux_miura3_stencil_01(): + mesh = SimpleMesh() + + z_lsq_coeff_1 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_2 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_3 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_4 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_5 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_6 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_7 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_8 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_9 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_10 = random_field(mesh, CellDim, KDim) + z_quad_vector_sum_1 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_2 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_3 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_4 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_5 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_6 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_7 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_8 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_9 = random_field(mesh, EdgeDim, KDim) + z_quad_vector_sum_10 = random_field(mesh, EdgeDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + z_dreg_area = random_field(mesh, EdgeDim, KDim) + cell_rel_idx_dsl = random_mask(mesh, EdgeDim, KDim, dtype=int32) + p_out_e_miura3 = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura3_stencil_01_numpy( + mesh.e2c, + np.asarray(z_lsq_coeff_1), + np.asarray(z_lsq_coeff_2), + np.asarray(z_lsq_coeff_3), + np.asarray(z_lsq_coeff_4), + np.asarray(z_lsq_coeff_5), + np.asarray(z_lsq_coeff_6), + np.asarray(z_lsq_coeff_7), + np.asarray(z_lsq_coeff_8), + np.asarray(z_lsq_coeff_9), + np.asarray(z_lsq_coeff_10), + np.asarray(z_quad_vector_sum_1), + np.asarray(z_quad_vector_sum_2), + np.asarray(z_quad_vector_sum_3), + np.asarray(z_quad_vector_sum_4), + np.asarray(z_quad_vector_sum_5), + np.asarray(z_quad_vector_sum_6), + np.asarray(z_quad_vector_sum_7), + np.asarray(z_quad_vector_sum_8), + np.asarray(z_quad_vector_sum_9), + np.asarray(z_quad_vector_sum_10), + np.asarray(z_dreg_area), + np.asarray(p_mass_flx_e), + np.asarray(cell_rel_idx_dsl), + ) + + upwind_hflux_miura3_stencil_01( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + z_lsq_coeff_4, + z_lsq_coeff_5, + z_lsq_coeff_6, + z_lsq_coeff_7, + z_lsq_coeff_8, + z_lsq_coeff_9, + z_lsq_coeff_10, + z_quad_vector_sum_1, + z_quad_vector_sum_2, + z_quad_vector_sum_3, + z_quad_vector_sum_4, + z_quad_vector_sum_5, + z_quad_vector_sum_6, + z_quad_vector_sum_7, + z_quad_vector_sum_8, + z_quad_vector_sum_9, + z_quad_vector_sum_10, + z_dreg_area, + p_mass_flx_e, + cell_rel_idx_dsl, + p_out_e_miura3, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(np.asarray(p_out_e_miura3), ref) diff --git a/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_01.py new file mode 100644 index 000000000..e7460b7e5 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_01.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 numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_01 import ( + upwind_hflux_miura_cycl_stencil_01, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura_cycl_stencil_01_numpy( + e2c: np.array, + z_lsq_coeff_1_dsl: np.array, + z_lsq_coeff_2_dsl: np.array, + z_lsq_coeff_3_dsl: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: int32, +): + z_lsq_coeff_1_dsl_e2c = z_lsq_coeff_1_dsl[e2c] + z_lsq_coeff_2_dsl_e2c = z_lsq_coeff_2_dsl[e2c] + z_lsq_coeff_3_dsl_e2c = z_lsq_coeff_3_dsl[e2c] + + z_tracer_mflx_dsl = ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_dsl_e2c[:, 1], + z_lsq_coeff_1_dsl_e2c[:, 0], + ) + + distv_bary_1 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_dsl_e2c[:, 1], + z_lsq_coeff_2_dsl_e2c[:, 0], + ) + + distv_bary_2 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_dsl_e2c[:, 1], + z_lsq_coeff_3_dsl_e2c[:, 0], + ) + ) * p_mass_flx_e + + return z_tracer_mflx_dsl + + +def test_upwind_hflux_miura_cycl_stencil_01(): + mesh = SimpleMesh() + + z_lsq_coeff_1_dsl = random_field(mesh, CellDim, KDim) + z_lsq_coeff_2_dsl = random_field(mesh, CellDim, KDim) + z_lsq_coeff_3_dsl = random_field(mesh, CellDim, KDim) + distv_bary_1 = random_field(mesh, EdgeDim, KDim) + distv_bary_2 = random_field(mesh, EdgeDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + cell_rel_idx_dsl = random_mask(mesh, EdgeDim, KDim, dtype=int32) + z_tracer_mflx_dsl = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura_cycl_stencil_01_numpy( + mesh.e2c, + np.asarray(z_lsq_coeff_1_dsl), + np.asarray(z_lsq_coeff_2_dsl), + np.asarray(z_lsq_coeff_3_dsl), + np.asarray(distv_bary_1), + np.asarray(distv_bary_2), + np.asarray(p_mass_flx_e), + np.asarray(cell_rel_idx_dsl), + ) + + upwind_hflux_miura_cycl_stencil_01( + z_lsq_coeff_1_dsl, + z_lsq_coeff_2_dsl, + z_lsq_coeff_3_dsl, + distv_bary_1, + distv_bary_2, + p_mass_flx_e, + cell_rel_idx_dsl, + z_tracer_mflx_dsl, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(ref, z_tracer_mflx_dsl) diff --git a/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_02.py new file mode 100644 index 000000000..823b0e046 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_02.py @@ -0,0 +1,101 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 +from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_02 import ( + upwind_hflux_miura_cycl_stencil_02, +) +from icon4py.model.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura_cycl_stencil_02_numpy( + c2e: np.array, + nsub: int32, + p_mass_flx_e: np.array, + geofac_div: np.array, + z_rhofluxdiv_c: np.array, + z_tracer_mflx: np.array, + z_rho_now: np.array, + z_tracer_now: np.array, + z_dtsub: float, +): + p_mass_flx_e_c2e = p_mass_flx_e[c2e] + geofac_div = np.expand_dims(geofac_div, axis=-1) + z_tracer_mflx_c2e = z_tracer_mflx[c2e] + + z_rhofluxdiv_c_out = ( + np.sum(p_mass_flx_e_c2e * geofac_div, axis=1) if nsub == int32(1) else z_rhofluxdiv_c + ) + z_fluxdiv_c_dsl = np.sum(z_tracer_mflx_c2e * geofac_div, axis=1) + + z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out + + z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl + + return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) + + +def test_upwind_hflux_miura_cycl_stencil_02(): + mesh = SimpleMesh() + nsub = int32(1) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + geofac_div = random_field(mesh, CellDim, C2EDim) + z_rhofluxdiv_c = random_field(mesh, CellDim, KDim) + z_tracer_mflx = random_field(mesh, EdgeDim, KDim) + z_rho_now = random_field(mesh, CellDim, KDim) + z_tracer_now = random_field(mesh, CellDim, KDim) + z_dtsub = 0.5 + z_rhofluxdiv_c_out = random_field(mesh, CellDim, KDim) + z_fluxdiv_c_dsl = random_field(mesh, CellDim, KDim) + z_rho_new_dsl = random_field(mesh, CellDim, KDim) + z_tracer_new_dsl = random_field(mesh, CellDim, KDim) + + ref_1, ref_2, ref_3, ref_4 = upwind_hflux_miura_cycl_stencil_02_numpy( + mesh.c2e, + nsub, + np.asarray(p_mass_flx_e), + np.asarray(geofac_div), + np.asarray(z_rhofluxdiv_c), + np.asarray(z_tracer_mflx), + np.asarray(z_rho_now), + np.asarray(z_tracer_now), + z_dtsub, + ) + + upwind_hflux_miura_cycl_stencil_02( + nsub, + p_mass_flx_e, + geofac_div, + z_rhofluxdiv_c, + z_tracer_mflx, + z_rho_now, + z_tracer_now, + z_dtsub, + z_rhofluxdiv_c_out, + z_fluxdiv_c_dsl, + z_rho_new_dsl, + z_tracer_new_dsl, + offset_provider={ + "C2CE": StridedNeighborOffsetProvider(CellDim, CEDim, mesh.n_c2e), + "C2E": mesh.get_c2e_offset_provider(), + }, + ) + assert np.allclose(ref_1, z_rhofluxdiv_c_out) + assert np.allclose(ref_2, z_fluxdiv_c_dsl) + assert np.allclose(ref_3, z_rho_new_dsl) + assert np.allclose(ref_4, z_tracer_new_dsl) diff --git a/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03a.py new file mode 100644 index 000000000..6e55efa77 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03a.py @@ -0,0 +1,49 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03a import ( + upwind_hflux_miura_cycl_stencil_03a, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura_cycl_stencil_03a_numpy( + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, +): + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float(2) + return p_out_e + + +def test_upwind_hflux_miura_cycl_stencil_03a(): + mesh = SimpleMesh() + z_tracer_mflx_1_dsl = random_field(mesh, EdgeDim, KDim) + z_tracer_mflx_2_dsl = random_field(mesh, EdgeDim, KDim) + p_out_e = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura_cycl_stencil_03a_numpy( + np.asarray(z_tracer_mflx_1_dsl), + np.asarray(z_tracer_mflx_2_dsl), + ) + + upwind_hflux_miura_cycl_stencil_03a( + z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl, + p_out_e, + offset_provider={}, + ) + assert np.allclose(ref, p_out_e) diff --git a/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03b.py new file mode 100644 index 000000000..91470d8eb --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_cycl_stencil_03b.py @@ -0,0 +1,53 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_cycl_stencil_03b import ( + upwind_hflux_miura_cycl_stencil_03b, +) +from icon4py.model.common.dimension import EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura_cycl_stencil_03b_numpy( + z_tracer_mflx_1_dsl: np.array, + z_tracer_mflx_2_dsl: np.array, + z_tracer_mflx_3_dsl: np.array, +): + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) + return p_out_e + + +def test_upwind_hflux_miura_cycl_stencil_03b(): + mesh = SimpleMesh() + z_tracer_mflx_1_dsl = random_field(mesh, EdgeDim, KDim) + z_tracer_mflx_2_dsl = random_field(mesh, EdgeDim, KDim) + z_tracer_mflx_3_dsl = random_field(mesh, EdgeDim, KDim) + p_out_e = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura_cycl_stencil_03b_numpy( + np.asarray(z_tracer_mflx_1_dsl), + np.asarray(z_tracer_mflx_2_dsl), + np.asarray(z_tracer_mflx_3_dsl), + ) + + upwind_hflux_miura_cycl_stencil_03b( + z_tracer_mflx_1_dsl, + z_tracer_mflx_2_dsl, + z_tracer_mflx_3_dsl, + p_out_e, + offset_provider={}, + ) + assert np.allclose(ref, p_out_e) diff --git a/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_stencil_01.py new file mode 100644 index 000000000..78a0e595c --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_upwind_hflux_miura_stencil_01.py @@ -0,0 +1,99 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.upwind_hflux_miura_stencil_01 import ( + upwind_hflux_miura_stencil_01, +) +from icon4py.model.common.dimension import CellDim, EdgeDim, KDim +from icon4py.model.common.test_utils.helpers import constant_field, random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def upwind_hflux_miura_stencil_01_numpy( + e2c: np.array, + z_lsq_coeff_1: np.array, + z_lsq_coeff_2: np.array, + z_lsq_coeff_3: np.array, + distv_bary_1: np.array, + distv_bary_2: np.array, + p_mass_flx_e: np.array, + cell_rel_idx_dsl: np.array, +) -> np.array: + + z_lsq_coeff_1_e2c = z_lsq_coeff_1[e2c] + z_lsq_coeff_2_e2c = z_lsq_coeff_2[e2c] + z_lsq_coeff_3_e2c = z_lsq_coeff_3[e2c] + + p_out_e = ( + np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_e2c[:, 1], + z_lsq_coeff_1_e2c[:, 0], + ) + + distv_bary_1 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_e2c[:, 1], + z_lsq_coeff_2_e2c[:, 0], + ) + + distv_bary_2 + * np.where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_e2c[:, 1], + z_lsq_coeff_3_e2c[:, 0], + ) + ) * p_mass_flx_e + + return p_out_e + + +def test_upwind_hflux_miura_stencil_01(): + mesh = SimpleMesh() + + z_lsq_coeff_1 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_2 = random_field(mesh, CellDim, KDim) + z_lsq_coeff_3 = random_field(mesh, CellDim, KDim) + distv_bary_1 = random_field(mesh, EdgeDim, KDim) + distv_bary_2 = random_field(mesh, EdgeDim, KDim) + p_mass_flx_e = random_field(mesh, EdgeDim, KDim) + cell_rel_idx_dsl = constant_field(mesh, 0, EdgeDim, KDim, dtype=int32) + p_out_e = zero_field(mesh, EdgeDim, KDim) + + ref = upwind_hflux_miura_stencil_01_numpy( + mesh.e2c, + np.asarray(z_lsq_coeff_1), + np.asarray(z_lsq_coeff_2), + np.asarray(z_lsq_coeff_3), + np.asarray(distv_bary_1), + np.asarray(distv_bary_2), + np.asarray(p_mass_flx_e), + np.asarray(cell_rel_idx_dsl), + ) + + upwind_hflux_miura_stencil_01( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + distv_bary_1, + distv_bary_2, + p_mass_flx_e, + cell_rel_idx_dsl, + p_out_e, + offset_provider={ + "E2C": mesh.get_e2c_offset_provider(), + }, + ) + assert np.allclose(p_out_e, ref) diff --git a/model/atmosphere/advection/advection_tests/test_upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/advection_tests/test_upwind_vflux_ppm_stencil_01.py new file mode 100644 index 000000000..e66c64f36 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_upwind_vflux_ppm_stencil_01.py @@ -0,0 +1,47 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.upwind_vflux_ppm_stencil_01 import ( + upwind_vflux_ppm_stencil_01, +) +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def upwind_vflux_ppm_stencil_01_numpy( + z_face_up: np.ndarray, z_face_low: np.ndarray, p_cc: np.ndarray +) -> tuple[np.ndarray]: + z_delta_q = 0.5 * (z_face_up - z_face_low) + z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) + return z_delta_q, z_a1 + + +def test_upwind_vflux_ppm_stencil_01(): + mesh = SimpleMesh() + z_face_up = random_field(mesh, CellDim, KDim) + z_face_down = random_field(mesh, CellDim, KDim) + p_cc = random_field(mesh, CellDim, KDim) + z_delta_q = zero_field(mesh, CellDim, KDim) + z_a1 = zero_field(mesh, CellDim, KDim) + + ref_z_delta_q, ref_z_a1 = upwind_vflux_ppm_stencil_01_numpy( + np.asarray(z_face_up), np.asarray(z_face_down), np.asarray(p_cc) + ) + + upwind_vflux_ppm_stencil_01(z_face_up, z_face_down, p_cc, z_delta_q, z_a1, offset_provider={}) + + assert np.allclose(ref_z_delta_q, z_delta_q) + assert np.allclose(ref_z_a1, z_a1) diff --git a/model/atmosphere/advection/advection_tests/test_vert_adv_stencil_01.py b/model/atmosphere/advection/advection_tests/test_vert_adv_stencil_01.py new file mode 100644 index 000000000..15e3f247e --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_vert_adv_stencil_01.py @@ -0,0 +1,73 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np + +from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def vert_adv_stencil_01_numpy( + tracer_now: np.array, + rhodz_now: np.array, + p_mflx_tracer_v: np.array, + deepatmo_divzl: np.array, + deepatmo_divzu: np.array, + rhodz_new: np.array, + p_dtime, +) -> np.array: + + tracer_new = ( + tracer_now * rhodz_now + + p_dtime + * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu) + ) / rhodz_new + + return tracer_new + + +def test_vert_adv_stencil_01(): + mesh = SimpleMesh() + + tracer_now = random_field(mesh, CellDim, KDim) + rhodz_now = random_field(mesh, CellDim, KDim) + p_mflx_tracer_v = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + deepatmo_divzl = random_field(mesh, KDim) + deepatmo_divzu = random_field(mesh, KDim) + rhodz_new = random_field(mesh, CellDim, KDim) + tracer_new = zero_field(mesh, CellDim, KDim) + p_dtime = np.float64(5.0) + + ref = vert_adv_stencil_01_numpy( + np.asarray(tracer_now), + np.asarray(rhodz_now), + np.asarray(p_mflx_tracer_v), + np.asarray(deepatmo_divzl), + np.asarray(deepatmo_divzu), + np.asarray(rhodz_new), + p_dtime, + ) + vert_adv_stencil_01( + tracer_now, + rhodz_now, + p_mflx_tracer_v, + deepatmo_divzl, + deepatmo_divzu, + rhodz_new, + tracer_new, + p_dtime, + offset_provider={"Koff": KDim}, + ) + assert np.allclose(tracer_new[:, :-1], ref[:, :-1]) diff --git a/model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_01.py b/model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_01.py new file mode 100644 index 000000000..d199edec9 --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_01.py @@ -0,0 +1,54 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_01 import v_limit_prbl_sm_stencil_01 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def v_limit_prbl_sm_stencil_01_numpy( + p_face: np.array, + p_cc: np.array, +): + + z_delta = p_face[:, :-1] - p_face[:, 1:] + z_a6i = 6.0 * (p_cc - 0.5 * (p_face[:, :-1] + p_face[:, 1:])) + + l_limit = np.where(np.abs(z_delta) < -1 * z_a6i, int32(1), int32(0)) + + return l_limit + + +def test_v_limit_prbl_sm_stencil_01(): + mesh = SimpleMesh() + p_cc = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + l_limit = zero_field(mesh, CellDim, KDim, dtype=int32) + + l_limit_ref = v_limit_prbl_sm_stencil_01_numpy( + np.asarray(p_face), + np.asarray(p_cc), + ) + + v_limit_prbl_sm_stencil_01( + p_face, + p_cc, + l_limit, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(l_limit_ref, l_limit) diff --git a/model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_02.py b/model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_02.py new file mode 100644 index 000000000..9c58230fd --- /dev/null +++ b/model/atmosphere/advection/advection_tests/test_vlimit_prbl_sm_stencil_02.py @@ -0,0 +1,70 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import numpy as np +from gt4py.next.ffront.fbuiltins import int32 + +from icon4py.model.atmosphere.advection.v_limit_prbl_sm_stencil_02 import v_limit_prbl_sm_stencil_02 +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field, random_mask, zero_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def v_limit_prbl_sm_stencil_02_numpy( + l_limit: np.array, + p_face: np.array, + p_cc: np.array, +): + + q_face_up, q_face_low = np.where( + l_limit != int32(0), + np.where( + (p_cc < np.minimum(p_face[:, :-1], p_face[:, 1:])), + (p_cc, p_cc), + np.where( + p_face[:, :-1] > p_face[:, 1:], + (3.0 * p_cc - 2.0 * p_face[:, 1:], p_face[:, 1:]), + (p_face[:, :-1], 3.0 * p_cc - 2.0 * p_face[:, :-1]), + ), + ), + (p_face[:, :-1], p_face[:, 1:]), + ) + + return q_face_up, q_face_low + + +def test_v_limit_prbl_sm_stencil_02(): + mesh = SimpleMesh() + l_limit = random_mask(mesh, CellDim, KDim, dtype=int32) + p_cc = random_field(mesh, CellDim, KDim) + p_face = random_field(mesh, CellDim, KDim, extend={KDim: 1}) + p_face_up = zero_field(mesh, CellDim, KDim) + p_face_low = zero_field(mesh, CellDim, KDim) + + p_face_up_ref, p_face_low_ref = v_limit_prbl_sm_stencil_02_numpy( + np.asarray(l_limit), + np.asarray(p_face), + np.asarray(p_cc), + ) + + v_limit_prbl_sm_stencil_02( + l_limit, + p_face, + p_cc, + p_face_up, + p_face_low, + offset_provider={"Koff": KDim}, + ) + + assert np.allclose(p_face_up_ref[:, :-1], p_face_up[:, :-1]) + assert np.allclose(p_face_low_ref[:, :-1], p_face_low[:, :-1]) diff --git a/model/atmosphere/advection/pyproject.toml b/model/atmosphere/advection/pyproject.toml new file mode 100644 index 000000000..f26f25f3b --- /dev/null +++ b/model/atmosphere/advection/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" +] +description = "ICON advection." +dynamic = ['version'] +license = {file = "LICENSE"} +name = "icon4py-atmosphere-advection" +readme = "README.md" +requires-python = ">=3.10" + +[project.urls] +repository = "https://github.com/C2SM/icon4py" + +[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 = ['advection'] + +[tool.isort] +force_grid_wrap = 0 +include_trailing_comma = true +known_first_party = ['icon4py'] +known_third_party = ['gt4py'] +lexicographical = true +line_length = 100 # It should be the same as in `tool.black.line-length` above +lines_after_imports = 2 +# known_tests = ['tests_utils'] +multi_line_output = 3 +profile = 'black' +sections = ['FUTURE', 'STDLIB', 'THIRDPARTY', 'FIRSTPARTY', '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] +markers = 'slow_tests: marks tests as slow' +testpaths = ['tests', 'advection_tests'] + +[tool.setuptools.dynamic] +version = {attr = 'icon4py.model.atmosphere.advection.__init__.__version__'} + +[tool.setuptools.package-data] +'icon4py.model.atmosphere.advection' = ['py.typed'] diff --git a/model/atmosphere/advection/requirements-dev.txt b/model/atmosphere/advection/requirements-dev.txt new file mode 100644 index 000000000..32d385e08 --- /dev/null +++ b/model/atmosphere/advection/requirements-dev.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements-dev.txt +-e ../../common +-e . diff --git a/model/atmosphere/advection/requirements.txt b/model/atmosphere/advection/requirements.txt new file mode 100644 index 000000000..79787ec13 --- /dev/null +++ b/model/atmosphere/advection/requirements.txt @@ -0,0 +1,3 @@ +-r ../../../base-requirements.txt +../../common +. diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__init__.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__init__.py new file mode 100644 index 000000000..dab708955 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/__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/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py new file mode 100644 index 000000000..17eea2c69 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_01.py @@ -0,0 +1,41 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, where + +from icon4py.model.common.dimension import EdgeDim, KDim + + +@field_operator +def _btraj_dreg_stencil_01( + lcounterclock: bool, + p_vn: Field[[EdgeDim, KDim], float], + tangent_orientation: Field[[EdgeDim], float], +) -> Field[[EdgeDim, KDim], bool]: + + tangent_orientation = broadcast(tangent_orientation, (EdgeDim, KDim)) + lvn_sys_pos_true = where(p_vn * tangent_orientation >= 0.0, True, False) + mask_lcounterclock = broadcast(lcounterclock, (EdgeDim, KDim)) + lvn_sys_pos = where(mask_lcounterclock, lvn_sys_pos_true, False) + return lvn_sys_pos + + +@program +def btraj_dreg_stencil_01( + lcounterclock: bool, + p_vn: Field[[EdgeDim, KDim], float], + tangent_orientation: Field[[EdgeDim], float], + lvn_sys_pos: Field[[EdgeDim, KDim], bool], +): + _btraj_dreg_stencil_01(lcounterclock, p_vn, tangent_orientation, out=lvn_sys_pos) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py new file mode 100644 index 000000000..5984b77c5 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_02.py @@ -0,0 +1,46 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, sqrt, where + +from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim + + +@field_operator +def _btraj_dreg_stencil_02( + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + edge_cell_length: Field[[ECDim], float], + p_dt: float, +) -> Field[[EdgeDim, KDim], int32]: + + lvn_pos = where(p_vn >= 0.0, True, False) + traj_length = sqrt(p_vn * p_vn + p_vt * p_vt) * p_dt + e2c_length = where(lvn_pos, edge_cell_length(E2EC[0]), edge_cell_length(E2EC[1])) + opt_famask_dsl = where( + traj_length > 1.25 * broadcast(e2c_length, (EdgeDim, KDim)), int32(1), int32(0) + ) + + return opt_famask_dsl + + +@program +def btraj_dreg_stencil_02( + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + edge_cell_length: Field[[ECDim], float], + p_dt: float, + opt_famask_dsl: Field[[EdgeDim, KDim], int32], +): + _btraj_dreg_stencil_02(p_vn, p_vt, edge_cell_length, p_dt, out=opt_famask_dsl) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py new file mode 100644 index 000000000..6dce27f72 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/btraj_dreg_stencil_03.py @@ -0,0 +1,169 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32, where + +from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim + + +@field_operator +def _btraj_dreg_stencil_03( + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + cell_idx: Field[[ECDim], int32], + cell_blk: Field[[ECDim], int32], + edge_verts_1_x: Field[[EdgeDim], float], + edge_verts_2_x: Field[[EdgeDim], float], + edge_verts_1_y: Field[[EdgeDim], float], + edge_verts_2_y: Field[[EdgeDim], float], + pos_on_tplane_e_1_x: Field[[EdgeDim], float], + pos_on_tplane_e_2_x: Field[[EdgeDim], float], + pos_on_tplane_e_1_y: Field[[EdgeDim], float], + pos_on_tplane_e_2_y: Field[[EdgeDim], float], + primal_normal_cell_x: Field[[ECDim], float], + primal_normal_cell_y: Field[[ECDim], float], + dual_normal_cell_x: Field[[ECDim], float], + dual_normal_cell_y: Field[[ECDim], float], + lvn_sys_pos: Field[[EdgeDim, KDim], bool], + p_dt: float, +) -> tuple[ + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], +]: + lvn_pos = where(p_vn >= 0.0, True, False) + + p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) + p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) + p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) + + depart_pts_1_x = edge_verts_1_x - p_vn * p_dt + depart_pts_1_y = edge_verts_1_y - p_vt * p_dt + depart_pts_2_x = edge_verts_2_x - p_vn * p_dt + depart_pts_2_y = edge_verts_2_y - p_vt * p_dt + + pos_on_tplane_e_x = where(lvn_pos, pos_on_tplane_e_1_x, pos_on_tplane_e_2_x) + pos_on_tplane_e_y = where(lvn_pos, pos_on_tplane_e_1_y, pos_on_tplane_e_2_y) + + pos_dreg_vert_c_1_x = edge_verts_1_x - pos_on_tplane_e_x + pos_dreg_vert_c_1_y = edge_verts_1_y - pos_on_tplane_e_y + pos_dreg_vert_c_2_x = where(lvn_sys_pos, depart_pts_1_x, edge_verts_2_x) - pos_on_tplane_e_x + pos_dreg_vert_c_2_y = where(lvn_sys_pos, depart_pts_1_y, edge_verts_2_y) - pos_on_tplane_e_y + pos_dreg_vert_c_3_x = depart_pts_2_x - pos_on_tplane_e_x + pos_dreg_vert_c_3_y = depart_pts_2_y - pos_on_tplane_e_y + pos_dreg_vert_c_4_x = where(lvn_sys_pos, edge_verts_2_x, depart_pts_1_x) - pos_on_tplane_e_x + pos_dreg_vert_c_4_y = where(lvn_sys_pos, edge_verts_2_y, depart_pts_1_y) - pos_on_tplane_e_y + + pn_cell_1 = where(lvn_pos, primal_normal_cell_x(E2EC[0]), primal_normal_cell_x(E2EC[1])) + pn_cell_2 = where(lvn_pos, primal_normal_cell_y(E2EC[0]), primal_normal_cell_y(E2EC[1])) + dn_cell_1 = where(lvn_pos, dual_normal_cell_x(E2EC[0]), dual_normal_cell_x(E2EC[1])) + dn_cell_2 = where(lvn_pos, dual_normal_cell_y(E2EC[0]), dual_normal_cell_y(E2EC[1])) + + p_coords_dreg_v_1_lon_dsl = pos_dreg_vert_c_1_x * pn_cell_1 + pos_dreg_vert_c_1_y * dn_cell_1 + p_coords_dreg_v_2_lon_dsl = pos_dreg_vert_c_2_x * pn_cell_1 + pos_dreg_vert_c_2_y * dn_cell_1 + p_coords_dreg_v_3_lon_dsl = pos_dreg_vert_c_3_x * pn_cell_1 + pos_dreg_vert_c_3_y * dn_cell_1 + p_coords_dreg_v_4_lon_dsl = pos_dreg_vert_c_4_x * pn_cell_1 + pos_dreg_vert_c_4_y * dn_cell_1 + p_coords_dreg_v_1_lat_dsl = pos_dreg_vert_c_1_x * pn_cell_2 + pos_dreg_vert_c_1_y * dn_cell_2 + p_coords_dreg_v_2_lat_dsl = pos_dreg_vert_c_2_x * pn_cell_2 + pos_dreg_vert_c_2_y * dn_cell_2 + p_coords_dreg_v_3_lat_dsl = pos_dreg_vert_c_3_x * pn_cell_2 + pos_dreg_vert_c_3_y * dn_cell_2 + p_coords_dreg_v_4_lat_dsl = pos_dreg_vert_c_4_x * pn_cell_2 + pos_dreg_vert_c_4_y * dn_cell_2 + + return ( + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl, + ) + + +@program +def btraj_dreg_stencil_03( + p_vn: Field[[EdgeDim, KDim], float], + p_vt: Field[[EdgeDim, KDim], float], + cell_idx: Field[[ECDim], int32], + cell_blk: Field[[ECDim], int32], + edge_verts_1_x: Field[[EdgeDim], float], + edge_verts_2_x: Field[[EdgeDim], float], + edge_verts_1_y: Field[[EdgeDim], float], + edge_verts_2_y: Field[[EdgeDim], float], + pos_on_tplane_e_1_x: Field[[EdgeDim], float], + pos_on_tplane_e_2_x: Field[[EdgeDim], float], + pos_on_tplane_e_1_y: Field[[EdgeDim], float], + pos_on_tplane_e_2_y: Field[[EdgeDim], float], + primal_normal_cell_x: Field[[ECDim], float], + primal_normal_cell_y: Field[[ECDim], float], + dual_normal_cell_x: Field[[ECDim], float], + dual_normal_cell_y: Field[[ECDim], float], + lvn_sys_pos: Field[[EdgeDim, KDim], bool], + p_dt: float, + p_cell_idx: Field[[EdgeDim, KDim], int32], + p_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], + p_cell_blk: Field[[EdgeDim, KDim], int32], + p_coords_dreg_v_1_lon_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_lon_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_lon_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_lon_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_lat_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_lat_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_lat_dsl: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_lat_dsl: Field[[EdgeDim, KDim], float], +): + _btraj_dreg_stencil_03( + p_vn, + p_vt, + cell_idx, + cell_blk, + edge_verts_1_x, + edge_verts_2_x, + edge_verts_1_y, + edge_verts_2_y, + pos_on_tplane_e_1_x, + pos_on_tplane_e_2_x, + pos_on_tplane_e_1_y, + pos_on_tplane_e_2_y, + primal_normal_cell_x, + primal_normal_cell_y, + dual_normal_cell_x, + dual_normal_cell_y, + lvn_sys_pos, + p_dt, + out=( + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_coords_dreg_v_1_lon_dsl, + p_coords_dreg_v_2_lon_dsl, + p_coords_dreg_v_3_lon_dsl, + p_coords_dreg_v_4_lon_dsl, + p_coords_dreg_v_1_lat_dsl, + p_coords_dreg_v_2_lat_dsl, + p_coords_dreg_v_3_lat_dsl, + p_coords_dreg_v_4_lat_dsl, + ), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py new file mode 100644 index 000000000..a753bd238 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_01.py @@ -0,0 +1,744 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import sys + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + +from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim + + +sys.setrecursionlimit(5500) + + +# FUNCTIONS +# Checking turn when travelling along three points, used to check whether lines inters. +@field_operator +def ccw( + p0_lon: Field[[EdgeDim, KDim], float], + p0_lat: Field[[EdgeDim, KDim], float], + p1_lon: Field[[EdgeDim, KDim], float], + p1_lat: Field[[EdgeDim, KDim], float], + p2_lon: Field[[EdgeDim, KDim], float], + p2_lat: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], int32]: + + dx1 = p1_lon - p0_lon + dy1 = p1_lat - p0_lat + + dx2 = p2_lon - p0_lon + dy2 = p2_lat - p0_lat + + dx1dy2 = dx1 * dy2 + dy1dx2 = dy1 * dx2 + + lccw = where(dx1dy2 > dy1dx2, True, False) + ccw_out = where(lccw, 1, -1) # 1: clockwise, -1: counterclockwise + return ccw_out + + +# Checks whether two lines intersect +@field_operator +def lintersect( + line1_p1_lon: Field[[EdgeDim, KDim], float], + line1_p1_lat: Field[[EdgeDim, KDim], float], + line1_p2_lon: Field[[EdgeDim, KDim], float], + line1_p2_lat: Field[[EdgeDim, KDim], float], + line2_p1_lon: Field[[EdgeDim, KDim], float], + line2_p1_lat: Field[[EdgeDim, KDim], float], + line2_p2_lon: Field[[EdgeDim, KDim], float], + line2_p2_lat: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], bool]: + + intersect1 = ccw( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p1_lon, + line2_p1_lat, + ) * ccw( + line1_p1_lon, + line1_p1_lat, + line1_p2_lon, + line1_p2_lat, + line2_p2_lon, + line2_p2_lat, + ) + intersect2 = ccw( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p1_lon, + line1_p1_lat, + ) * ccw( + line2_p1_lon, + line2_p1_lat, + line2_p2_lon, + line2_p2_lat, + line1_p2_lon, + line1_p2_lat, + ) + lintersect_out = where((intersect1 + intersect2) == -2, True, False) + + return lintersect_out + + +# Compute intersection point of two lines in 2D +@field_operator +def line_intersect( + line1_p1_lon: Field[[EdgeDim, KDim], float], + line1_p1_lat: Field[[EdgeDim, KDim], float], + line1_p2_lon: Field[[EdgeDim, KDim], float], + line1_p2_lat: Field[[EdgeDim, KDim], float], + line2_p1_lon: Field[[EdgeDim, KDim], float], + line2_p1_lat: Field[[EdgeDim, KDim], float], + line2_p2_lon: Field[[EdgeDim, KDim], float], + line2_p2_lat: Field[[EdgeDim, KDim], float], +) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: + + m1 = (line1_p2_lat - line1_p1_lat) / (line1_p2_lon - line1_p1_lon) + m2 = (line2_p2_lat - line2_p1_lat) / (line2_p2_lon - line2_p1_lon) + + intersect_1 = (line2_p1_lat - line1_p1_lat + m1 * line1_p1_lon - m2 * line2_p1_lon) / (m1 - m2) + intersect_2 = line1_p1_lat + m1 * (intersect_1 - line1_p1_lon) + + return intersect_1, intersect_2 + + +@field_operator +def _divide_flux_area_list_stencil_01( + famask_int: Field[[EdgeDim, KDim], int32], + p_vn: Field[[EdgeDim, KDim], float], + ptr_v3_lon: Field[[ECDim], float], + ptr_v3_lat: Field[[ECDim], float], + tangent_orientation_dsl: Field[[EdgeDim], float], + dreg_patch0_1_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_1_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_2_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_2_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_3_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_3_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_4_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_4_lat_dsl: Field[[EdgeDim, KDim], float], +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], +]: + + arrival_pts_1_lon_dsl = dreg_patch0_1_lon_dsl + arrival_pts_1_lat_dsl = dreg_patch0_1_lat_dsl + arrival_pts_2_lon_dsl = dreg_patch0_2_lon_dsl + arrival_pts_2_lat_dsl = dreg_patch0_2_lat_dsl + depart_pts_1_lon_dsl = dreg_patch0_4_lon_dsl # indices have to be switched so that dep 1 belongs to arr 1 (and d2->a2) + depart_pts_1_lat_dsl = dreg_patch0_4_lat_dsl + depart_pts_2_lon_dsl = dreg_patch0_3_lon_dsl + depart_pts_2_lat_dsl = dreg_patch0_3_lat_dsl + + lvn_pos = where(p_vn >= 0.0, True, False) + + # get flux area departure-line segment + fl_line_p1_lon = depart_pts_1_lon_dsl + fl_line_p1_lat = depart_pts_1_lat_dsl + fl_line_p2_lon = depart_pts_2_lon_dsl + fl_line_p2_lat = depart_pts_2_lat_dsl + + # get triangle edge 1 (A1V3) + tri_line1_p1_lon = arrival_pts_1_lon_dsl + tri_line1_p1_lat = arrival_pts_1_lat_dsl + tri_line1_p2_lon = where( + lvn_pos, + broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), + broadcast(ptr_v3_lon(E2EC[1]), (EdgeDim, KDim)), + ) + tri_line1_p2_lat = where( + lvn_pos, + broadcast(ptr_v3_lat(E2EC[0]), (EdgeDim, KDim)), + broadcast(ptr_v3_lat(E2EC[1]), (EdgeDim, KDim)), + ) + + # get triangle edge 2 (A2V3) + tri_line2_p1_lon = arrival_pts_2_lon_dsl + tri_line2_p1_lat = arrival_pts_2_lat_dsl + tri_line2_p2_lon = where( + lvn_pos, + broadcast(ptr_v3_lon(E2EC[0]), (EdgeDim, KDim)), + broadcast(ptr_v3_lon(E2EC[1]), (EdgeDim, KDim)), + ) + tri_line2_p2_lat = where( + lvn_pos, + broadcast(ptr_v3_lat(E2EC[0]), (EdgeDim, KDim)), + broadcast(ptr_v3_lat(E2EC[1]), (EdgeDim, KDim)), + ) + + # Create first mask does departure-line segment intersects with A1V3 + lintersect_line1 = lintersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Create first mask does departure-line segment intersects with A2V3 + lintersect_line2 = lintersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + lvn_sys_pos = where( + (p_vn * broadcast(tangent_orientation_dsl, (EdgeDim, KDim))) >= 0.0, True, False + ) + famask_bool = where(famask_int == int32(1), True, False) + # ------------------------------------------------- Case 1 + mask_case1 = lintersect_line1 & lintersect_line2 & famask_bool + ps1_x, ps1_y = line_intersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + ps2_x, ps2_y = line_intersect( + fl_line_p1_lon, + fl_line_p1_lat, + fl_line_p2_lon, + fl_line_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + + # Case 1 - patch 0 + dreg_patch0_1_lon_dsl = where(mask_case1, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case1, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = where( + mask_case1, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = where( + mask_case1, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = where(mask_case1, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = where(mask_case1, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = where( + mask_case1, + where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = where( + mask_case1, + where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 1 - patch 1 + dreg_patch1_1_lon_vmask = where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_1_lat_vmask = where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_4_lon_vmask = where(mask_case1, arrival_pts_1_lon_dsl, 0.0) + dreg_patch1_4_lat_vmask = where(mask_case1, arrival_pts_1_lat_dsl, 0.0) + dreg_patch1_2_lon_vmask = where( + mask_case1, where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), 0.0 + ) + dreg_patch1_2_lat_vmask = where( + mask_case1, where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), 0.0 + ) + dreg_patch1_3_lon_vmask = where( + mask_case1, where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), 0.0 + ) + dreg_patch1_3_lat_vmask = where( + mask_case1, where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), 0.0 + ) + # Case 1 - patch 2 + dreg_patch2_1_lon_vmask = where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_1_lat_vmask = where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_4_lon_vmask = where(mask_case1, arrival_pts_2_lon_dsl, 0.0) + dreg_patch2_4_lat_vmask = where(mask_case1, arrival_pts_2_lat_dsl, 0.0) + dreg_patch2_2_lon_vmask = where( + mask_case1, where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), 0.0 + ) + dreg_patch2_2_lat_vmask = where( + mask_case1, where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), 0.0 + ) + dreg_patch2_3_lon_vmask = where( + mask_case1, where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), 0.0 + ) + dreg_patch2_3_lat_vmask = where( + mask_case1, where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), 0.0 + ) + + # ------------------------------------------------- Case 2a + mask_case2a = lintersect_line1 & (not lintersect_line2) & famask_bool + # Case 2a - patch 0 + dreg_patch0_1_lon_dsl = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = where( + mask_case2a, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, ps1_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = where( + mask_case2a, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, ps1_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = where(mask_case2a, depart_pts_2_lon_dsl, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = where(mask_case2a, depart_pts_2_lat_dsl, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = where( + mask_case2a, + where(lvn_sys_pos, ps1_x, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = where( + mask_case2a, + where(lvn_sys_pos, ps1_y, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 2a - patch 1 + dreg_patch1_1_lon_vmask = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_4_lon_vmask = where(mask_case2a, arrival_pts_1_lon_dsl, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = where(mask_case2a, arrival_pts_1_lat_dsl, dreg_patch1_4_lat_vmask) + dreg_patch1_2_lon_vmask = where( + mask_case2a, + where(lvn_sys_pos, ps1_x, depart_pts_1_lon_dsl), + dreg_patch1_2_lon_vmask, + ) + dreg_patch1_2_lat_vmask = where( + mask_case2a, + where(lvn_sys_pos, ps1_y, depart_pts_1_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = where( + mask_case2a, + where(lvn_sys_pos, depart_pts_1_lon_dsl, ps1_x), + dreg_patch1_3_lon_vmask, + ) + dreg_patch1_3_lat_vmask = where( + mask_case2a, + where(lvn_sys_pos, depart_pts_1_lat_dsl, ps1_y), + dreg_patch1_3_lat_vmask, + ) + # Case 2a - patch 2 + dreg_patch2_1_lon_vmask = where(mask_case2a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case2a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = where(mask_case2a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = where(mask_case2a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = where(mask_case2a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = where(mask_case2a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = where(mask_case2a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = where(mask_case2a, 0.0, dreg_patch2_4_lat_vmask) + + # -------------------------------------------------- Case 2b + mask_case2b = lintersect_line2 & (not lintersect_line1) & famask_bool + # Case 2b - patch 0 + dreg_patch0_1_lon_dsl = where(mask_case2b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case2b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = where( + mask_case2b, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = where( + mask_case2b, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = where(mask_case2b, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = where(mask_case2b, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = where( + mask_case2b, + where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = where( + mask_case2b, + where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 2b - patch 1 + dreg_patch1_1_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = where(mask_case2b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = where(mask_case2b, 0.0, dreg_patch1_4_lat_vmask) + # Case 2b - patch 2 + dreg_patch2_1_lon_vmask = where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_4_lon_vmask = where(mask_case2b, arrival_pts_2_lon_dsl, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = where(mask_case2b, arrival_pts_2_lat_dsl, dreg_patch2_4_lat_vmask) + dreg_patch2_2_lon_vmask = where( + mask_case2b, + where(lvn_sys_pos, depart_pts_2_lon_dsl, ps2_x), + dreg_patch2_2_lon_vmask, + ) + dreg_patch2_2_lat_vmask = where( + mask_case2b, + where(lvn_sys_pos, depart_pts_2_lat_dsl, ps2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = where( + mask_case2b, + where(lvn_sys_pos, ps2_x, depart_pts_2_lon_dsl), + dreg_patch2_3_lon_vmask, + ) + dreg_patch2_3_lat_vmask = where( + mask_case2b, + where(lvn_sys_pos, ps2_y, depart_pts_2_lat_dsl), + dreg_patch2_3_lat_vmask, + ) + + # flux area edge 1 and 2 + fl_e1_p1_lon = arrival_pts_1_lon_dsl + fl_e1_p1_lat = arrival_pts_1_lat_dsl + fl_e1_p2_lon = depart_pts_1_lon_dsl + fl_e1_p2_lat = depart_pts_1_lat_dsl + fl_e2_p1_lon = arrival_pts_2_lon_dsl + fl_e2_p1_lat = arrival_pts_2_lat_dsl + fl_e2_p2_lon = depart_pts_2_lon_dsl + fl_e2_p2_lat = depart_pts_2_lat_dsl + + # ----------------------------------------------- Case 3a + # Check whether flux area edge 2 intersects with triangle edge 1 + lintersect_e2_line1 = lintersect( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + mask_case3a = lintersect_e2_line1 & famask_bool + pi1_x, pi1_y = line_intersect( + fl_e2_p1_lon, + fl_e2_p1_lat, + fl_e2_p2_lon, + fl_e2_p2_lat, + tri_line1_p1_lon, + tri_line1_p1_lat, + tri_line1_p2_lon, + tri_line1_p2_lat, + ) + # Case 3a - patch 0 + dreg_patch0_1_lon_dsl = where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_2_lon_dsl = where( + mask_case3a, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, depart_pts_1_lon_dsl), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = where( + mask_case3a, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, depart_pts_1_lat_dsl), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = where(mask_case3a, ps2_x, dreg_patch0_3_lon_dsl) + dreg_patch0_3_lat_dsl = where(mask_case3a, ps2_y, dreg_patch0_3_lat_dsl) + dreg_patch0_4_lon_dsl = where( + mask_case3a, + where(lvn_sys_pos, depart_pts_1_lon_dsl, arrival_pts_2_lon_dsl), + dreg_patch0_4_lon_dsl, + ) + dreg_patch0_4_lat_dsl = where( + mask_case3a, + where(lvn_sys_pos, depart_pts_1_lat_dsl, arrival_pts_2_lat_dsl), + dreg_patch0_4_lat_dsl, + ) + # Case 3a - patch 1 + dreg_patch1_1_lon_vmask = where(mask_case3a, arrival_pts_1_lon_dsl, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case3a, arrival_pts_1_lat_dsl, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = where( + mask_case3a, + where(lvn_sys_pos, pi1_x, depart_pts_2_lon_dsl), + dreg_patch1_2_lon_vmask, + ) + dreg_patch1_2_lat_vmask = where( + mask_case3a, + where(lvn_sys_pos, pi1_y, depart_pts_2_lat_dsl), + dreg_patch1_2_lat_vmask, + ) + dreg_patch1_3_lon_vmask = where(mask_case3a, depart_pts_1_lon_dsl, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = where(mask_case3a, depart_pts_1_lat_dsl, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = where( + mask_case3a, + where(lvn_sys_pos, depart_pts_1_lon_dsl, pi1_x), + dreg_patch1_4_lon_vmask, + ) + dreg_patch1_4_lat_vmask = where( + mask_case3a, + where(lvn_sys_pos, depart_pts_1_lat_dsl, pi1_y), + dreg_patch1_4_lat_vmask, + ) + # Case 3a - patch 2 + dreg_patch2_1_lon_vmask = where(mask_case3a, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case3a, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = where(mask_case3a, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = where(mask_case3a, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = where(mask_case3a, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = where(mask_case3a, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = where(mask_case3a, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = where(mask_case3a, 0.0, dreg_patch2_4_lat_vmask) + + # ------------------------------------------------ Case 3b + # Check whether flux area edge 1 intersects with triangle edge 2 + lintersect_e1_line2 = lintersect( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + mask_case3b = lintersect_e1_line2 & famask_bool + pi2_x, pi2_y = line_intersect( + fl_e1_p1_lon, + fl_e1_p1_lat, + fl_e1_p2_lon, + fl_e1_p2_lat, + tri_line2_p1_lon, + tri_line2_p1_lat, + tri_line2_p2_lon, + tri_line2_p2_lat, + ) + # Case 3b - patch 0 + dreg_patch0_1_lon_dsl = where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_1_lon_dsl) + dreg_patch0_1_lat_dsl = where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_1_lat_dsl) + dreg_patch0_4_lon_dsl = where(mask_case3b, arrival_pts_1_lon_dsl, dreg_patch0_4_lon_dsl) + dreg_patch0_4_lat_dsl = where(mask_case3b, arrival_pts_1_lat_dsl, dreg_patch0_4_lat_dsl) + dreg_patch0_2_lon_dsl = where( + mask_case3b, + where(lvn_sys_pos, arrival_pts_2_lon_dsl, pi2_x), + dreg_patch0_2_lon_dsl, + ) + dreg_patch0_2_lat_dsl = where( + mask_case3b, + where(lvn_sys_pos, arrival_pts_2_lat_dsl, pi2_y), + dreg_patch0_2_lat_dsl, + ) + dreg_patch0_3_lon_dsl = where( + mask_case3b, + where(lvn_sys_pos, pi2_x, arrival_pts_2_lon_dsl), + dreg_patch0_3_lon_dsl, + ) + dreg_patch0_3_lat_dsl = where( + mask_case3b, + where(lvn_sys_pos, pi2_y, arrival_pts_2_lat_dsl), + dreg_patch0_3_lat_dsl, + ) + # Case 3b - patch 1 + dreg_patch1_1_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = where(mask_case3b, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = where(mask_case3b, 0.0, dreg_patch1_4_lat_vmask) + # Case 3b - patch 2 + dreg_patch2_1_lon_vmask = where(mask_case3b, arrival_pts_2_lon_dsl, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case3b, arrival_pts_2_lat_dsl, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = where( + mask_case3b, + where(lvn_sys_pos, depart_pts_2_lon_dsl, pi2_x), + dreg_patch2_2_lon_vmask, + ) + dreg_patch2_2_lat_vmask = where( + mask_case3b, + where(lvn_sys_pos, depart_pts_2_lat_dsl, pi2_y), + dreg_patch2_2_lat_vmask, + ) + dreg_patch2_3_lon_vmask = where(mask_case3b, depart_pts_1_lon_dsl, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = where(mask_case3b, depart_pts_1_lat_dsl, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = where( + mask_case3b, + where(lvn_sys_pos, pi2_x, depart_pts_2_lon_dsl), + dreg_patch2_4_lon_vmask, + ) + dreg_patch2_4_lat_vmask = where( + mask_case3b, + where(lvn_sys_pos, pi2_y, depart_pts_2_lat_dsl), + dreg_patch2_4_lat_vmask, + ) + + # --------------------------------------------- Case 4 + # NB: Next line acts as the "ELSE IF", indices that already previously matched one of the above conditions + # can't be overwritten by this new condition. + indices_previously_matched = mask_case3b | mask_case3a | mask_case2b | mask_case2a | mask_case1 + # mask_case4 = (abs(p_vn) < 0.1) & famask_bool & (not indices_previously_matched) we insert also the error indices + mask_case4 = famask_bool & (not indices_previously_matched) + # Case 4 - patch 0 - no change + # Case 4 - patch 1 + dreg_patch1_1_lon_vmask = where(mask_case4, 0.0, dreg_patch1_1_lon_vmask) + dreg_patch1_1_lat_vmask = where(mask_case4, 0.0, dreg_patch1_1_lat_vmask) + dreg_patch1_2_lon_vmask = where(mask_case4, 0.0, dreg_patch1_2_lon_vmask) + dreg_patch1_2_lat_vmask = where(mask_case4, 0.0, dreg_patch1_2_lat_vmask) + dreg_patch1_3_lon_vmask = where(mask_case4, 0.0, dreg_patch1_3_lon_vmask) + dreg_patch1_3_lat_vmask = where(mask_case4, 0.0, dreg_patch1_3_lat_vmask) + dreg_patch1_4_lon_vmask = where(mask_case4, 0.0, dreg_patch1_4_lon_vmask) + dreg_patch1_4_lat_vmask = where(mask_case4, 0.0, dreg_patch1_4_lat_vmask) + # Case 4 - patch 2 + dreg_patch2_1_lon_vmask = where(mask_case4, 0.0, dreg_patch2_1_lon_vmask) + dreg_patch2_1_lat_vmask = where(mask_case4, 0.0, dreg_patch2_1_lat_vmask) + dreg_patch2_2_lon_vmask = where(mask_case4, 0.0, dreg_patch2_2_lon_vmask) + dreg_patch2_2_lat_vmask = where(mask_case4, 0.0, dreg_patch2_2_lat_vmask) + dreg_patch2_3_lon_vmask = where(mask_case4, 0.0, dreg_patch2_3_lon_vmask) + dreg_patch2_3_lat_vmask = where(mask_case4, 0.0, dreg_patch2_3_lat_vmask) + dreg_patch2_4_lon_vmask = where(mask_case4, 0.0, dreg_patch2_4_lon_vmask) + dreg_patch2_4_lat_vmask = where(mask_case4, 0.0, dreg_patch2_4_lat_vmask) + + return ( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ) + + +@program +def divide_flux_area_list_stencil_01( + famask_int: Field[[EdgeDim, KDim], int32], + p_vn: Field[[EdgeDim, KDim], float], + ptr_v3_lon: Field[[ECDim], float], + ptr_v3_lat: Field[[ECDim], float], + tangent_orientation_dsl: Field[[EdgeDim], float], + dreg_patch0_1_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_1_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_2_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_2_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_3_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_3_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_4_lon_dsl: Field[[EdgeDim, KDim], float], + dreg_patch0_4_lat_dsl: Field[[EdgeDim, KDim], float], + dreg_patch1_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lat_vmask: Field[[EdgeDim, KDim], float], +): + _divide_flux_area_list_stencil_01( + famask_int, + p_vn, + ptr_v3_lon, + ptr_v3_lat, + tangent_orientation_dsl, + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + out=( + dreg_patch0_1_lon_dsl, + dreg_patch0_1_lat_dsl, + dreg_patch0_2_lon_dsl, + dreg_patch0_2_lat_dsl, + dreg_patch0_3_lon_dsl, + dreg_patch0_3_lat_dsl, + dreg_patch0_4_lon_dsl, + dreg_patch0_4_lat_dsl, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + ), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py new file mode 100644 index 000000000..08650dade --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/divide_flux_area_list_stencil_02.py @@ -0,0 +1,267 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +import sys + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + +from icon4py.model.common.dimension import E2EC, ECDim, EdgeDim, KDim + + +sys.setrecursionlimit(5500) + + +@field_operator +def _divide_flux_area_list_stencil_02( + famask_int: Field[[EdgeDim, KDim], int32], + p_vn: Field[[EdgeDim, KDim], float], + bf_cc_patch1_lon: Field[[ECDim], float], + bf_cc_patch1_lat: Field[[ECDim], float], + bf_cc_patch2_lon: Field[[ECDim], float], + bf_cc_patch2_lat: Field[[ECDim], float], + butterfly_idx_patch1_vnpos: Field[[EdgeDim], int32], + butterfly_idx_patch1_vnneg: Field[[EdgeDim], int32], + butterfly_blk_patch1_vnpos: Field[[EdgeDim], int32], + butterfly_blk_patch1_vnneg: Field[[EdgeDim], int32], + butterfly_idx_patch2_vnpos: Field[[EdgeDim], int32], + butterfly_idx_patch2_vnneg: Field[[EdgeDim], int32], + butterfly_blk_patch2_vnpos: Field[[EdgeDim], int32], + butterfly_blk_patch2_vnneg: Field[[EdgeDim], int32], + dreg_patch1_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lat_vmask: Field[[EdgeDim, KDim], float], +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], int32], + Field[[EdgeDim, KDim], int32], +]: + famask_bool = where(famask_int == int32(1), True, False) + lvn_pos = where(p_vn >= 0.0, True, False) + # Translation of patch 1 and patch 2 in system relative to respective cell + bf_cc_patch1_lon = where( + famask_bool, + where(lvn_pos, bf_cc_patch1_lon(E2EC[0]), bf_cc_patch1_lon(E2EC[1])), + 0.0, + ) + bf_cc_patch1_lat = where( + famask_bool, + where(lvn_pos, bf_cc_patch1_lat(E2EC[0]), bf_cc_patch1_lat(E2EC[1])), + 0.0, + ) + bf_cc_patch2_lon = where( + famask_bool, + where(lvn_pos, bf_cc_patch2_lon(E2EC[0]), bf_cc_patch2_lon(E2EC[1])), + 0.0, + ) + bf_cc_patch2_lat = where( + famask_bool, + where(lvn_pos, bf_cc_patch2_lat(E2EC[0]), bf_cc_patch2_lat(E2EC[1])), + 0.0, + ) + + # patch1 in translated system + dreg_patch1_1_lon_vmask = dreg_patch1_1_lon_vmask - bf_cc_patch1_lon + dreg_patch1_1_lat_vmask = dreg_patch1_1_lat_vmask - bf_cc_patch1_lat + dreg_patch1_2_lon_vmask = dreg_patch1_2_lon_vmask - bf_cc_patch1_lon + dreg_patch1_2_lat_vmask = dreg_patch1_2_lat_vmask - bf_cc_patch1_lat + dreg_patch1_3_lon_vmask = dreg_patch1_3_lon_vmask - bf_cc_patch1_lon + dreg_patch1_3_lat_vmask = dreg_patch1_3_lat_vmask - bf_cc_patch1_lat + dreg_patch1_4_lon_vmask = dreg_patch1_4_lon_vmask - bf_cc_patch1_lon + dreg_patch1_4_lat_vmask = dreg_patch1_4_lat_vmask - bf_cc_patch1_lat + # patch2 in translated system + dreg_patch2_1_lon_vmask = dreg_patch2_1_lon_vmask - bf_cc_patch2_lon + dreg_patch2_1_lat_vmask = dreg_patch2_1_lat_vmask - bf_cc_patch2_lat + dreg_patch2_2_lon_vmask = dreg_patch2_2_lon_vmask - bf_cc_patch2_lon + dreg_patch2_2_lat_vmask = dreg_patch2_2_lat_vmask - bf_cc_patch2_lat + dreg_patch2_3_lon_vmask = dreg_patch2_3_lon_vmask - bf_cc_patch2_lon + dreg_patch2_3_lat_vmask = dreg_patch2_3_lat_vmask - bf_cc_patch2_lat + dreg_patch2_4_lon_vmask = dreg_patch2_4_lon_vmask - bf_cc_patch2_lon + dreg_patch2_4_lat_vmask = dreg_patch2_4_lat_vmask - bf_cc_patch2_lat + + # Store global index of the underlying grid cell + # Adapt dimensions to fit ofr multiple levels + butterfly_idx_patch1_vnpos_3d = broadcast(butterfly_idx_patch1_vnpos, (EdgeDim, KDim)) + butterfly_idx_patch1_vnneg_3d = broadcast(butterfly_idx_patch1_vnneg, (EdgeDim, KDim)) + butterfly_idx_patch2_vnpos_3d = broadcast(butterfly_idx_patch2_vnpos, (EdgeDim, KDim)) + butterfly_idx_patch2_vnneg_3d = broadcast(butterfly_idx_patch2_vnneg, (EdgeDim, KDim)) + butterfly_blk_patch1_vnpos_3d = broadcast(butterfly_blk_patch1_vnpos, (EdgeDim, KDim)) + butterfly_blk_patch1_vnneg_3d = broadcast(butterfly_blk_patch1_vnneg, (EdgeDim, KDim)) + butterfly_blk_patch2_vnpos_3d = broadcast(butterfly_blk_patch2_vnpos, (EdgeDim, KDim)) + butterfly_blk_patch2_vnneg_3d = broadcast(butterfly_blk_patch2_vnneg, (EdgeDim, KDim)) + patch1_cell_idx_vmask = where( + famask_bool, + where(lvn_pos, butterfly_idx_patch1_vnpos_3d, butterfly_idx_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_idx_vmask = where( + famask_bool, + where(lvn_pos, butterfly_idx_patch2_vnpos_3d, butterfly_idx_patch2_vnneg_3d), + int32(0), + ) + patch1_cell_blk_vmask = where( + famask_bool, + where(lvn_pos, butterfly_blk_patch1_vnpos_3d, butterfly_blk_patch1_vnneg_3d), + int32(0), + ) + patch2_cell_blk_vmask = where( + famask_bool, + where(lvn_pos, butterfly_blk_patch2_vnpos_3d, butterfly_blk_patch2_vnneg_3d), + int32(0), + ) + + return ( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask, + patch1_cell_blk_vmask, + patch2_cell_idx_vmask, + patch2_cell_blk_vmask, + ) + + +@program +def divide_flux_area_list_stencil_02( + famask_int: Field[[EdgeDim, KDim], int32], + p_vn: Field[[EdgeDim, KDim], float], + bf_cc_patch1_lon: Field[[ECDim], float], + bf_cc_patch1_lat: Field[[ECDim], float], + bf_cc_patch2_lon: Field[[ECDim], float], + bf_cc_patch2_lat: Field[[ECDim], float], + butterfly_idx_patch1_vnpos: Field[[EdgeDim], int32], + butterfly_idx_patch1_vnneg: Field[[EdgeDim], int32], + butterfly_blk_patch1_vnpos: Field[[EdgeDim], int32], + butterfly_blk_patch1_vnneg: Field[[EdgeDim], int32], + butterfly_idx_patch2_vnpos: Field[[EdgeDim], int32], + butterfly_idx_patch2_vnneg: Field[[EdgeDim], int32], + butterfly_blk_patch2_vnpos: Field[[EdgeDim], int32], + butterfly_blk_patch2_vnneg: Field[[EdgeDim], int32], + dreg_patch1_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch1_4_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_1_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_2_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_3_lat_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lon_vmask: Field[[EdgeDim, KDim], float], + dreg_patch2_4_lat_vmask: Field[[EdgeDim, KDim], float], + patch1_cell_idx_vmask: Field[[EdgeDim, KDim], int32], + patch1_cell_blk_vmask: Field[[EdgeDim, KDim], int32], + patch2_cell_idx_vmask: Field[[EdgeDim, KDim], int32], + patch2_cell_blk_vmask: Field[[EdgeDim, KDim], int32], +): + _divide_flux_area_list_stencil_02( + famask_int, + p_vn, + bf_cc_patch1_lon, + bf_cc_patch1_lat, + bf_cc_patch2_lon, + bf_cc_patch2_lat, + butterfly_idx_patch1_vnpos, + butterfly_idx_patch1_vnneg, + butterfly_blk_patch1_vnpos, + butterfly_blk_patch1_vnneg, + butterfly_idx_patch2_vnpos, + butterfly_idx_patch2_vnneg, + butterfly_blk_patch2_vnpos, + butterfly_blk_patch2_vnneg, + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + out=( + dreg_patch1_1_lon_vmask, + dreg_patch1_1_lat_vmask, + dreg_patch1_2_lon_vmask, + dreg_patch1_2_lat_vmask, + dreg_patch1_3_lon_vmask, + dreg_patch1_3_lat_vmask, + dreg_patch1_4_lon_vmask, + dreg_patch1_4_lat_vmask, + dreg_patch2_1_lon_vmask, + dreg_patch2_1_lat_vmask, + dreg_patch2_2_lon_vmask, + dreg_patch2_2_lat_vmask, + dreg_patch2_3_lon_vmask, + dreg_patch2_3_lat_vmask, + dreg_patch2_4_lon_vmask, + dreg_patch2_4_lat_vmask, + patch1_cell_idx_vmask, + patch1_cell_blk_vmask, + patch2_cell_idx_vmask, + patch2_cell_blk_vmask, + ), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py new file mode 100644 index 000000000..de605d3dd --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_01.py @@ -0,0 +1,90 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + +from icon4py.model.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_01a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) + zfac = (p_cc(Koff[+1]) - p_cc) / (p_cellhgt_mc_now(Koff[+1]) + p_cellhgt_mc_now) + z_slope = ( + p_cellhgt_mc_now + / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[+1])) + ) * ( + (2.0 * p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + + (p_cellhgt_mc_now + 2.0 * p_cellhgt_mc_now(Koff[+1])) * zfac_m1 + ) + + return z_slope + + +@field_operator +def _face_val_ppm_stencil_01b( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + zfac_m1 = (p_cc - p_cc(Koff[-1])) / (p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1])) + zfac = (p_cc - p_cc) / (p_cellhgt_mc_now + p_cellhgt_mc_now) + z_slope = ( + p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + p_cellhgt_mc_now) + ) * ( + (2.0 * p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) * zfac + + (p_cellhgt_mc_now + 2.0 * p_cellhgt_mc_now) * zfac_m1 + ) + + return z_slope + + +@field_operator +def _face_val_ppm_stencil_01( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + vert_idx: Field[[KDim], int32], + elev: int32, +) -> Field[[CellDim, KDim], float]: + + vert_idx = broadcast(vert_idx, (CellDim, KDim)) + + z_slope = where( + vert_idx == elev, + _face_val_ppm_stencil_01b(p_cc, p_cellhgt_mc_now), + _face_val_ppm_stencil_01a(p_cc, p_cellhgt_mc_now), + ) + + return z_slope + + +@program +def face_val_ppm_stencil_01( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + vert_idx: Field[[KDim], int32], + elev: int32, + z_slope: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_01( + p_cc, + p_cellhgt_mc_now, + vert_idx, + elev, + out=z_slope, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py new file mode 100644 index 000000000..1fc1cf235 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02.py @@ -0,0 +1,100 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where + +from icon4py.model.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_02a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc * (1.0 - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + ( + p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) + ) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1])) + + return p_face + + +@field_operator +def _face_val_ppm_stencil_02b( + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc + return p_face + + +@field_operator +def _face_val_ppm_stencil_02c( + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc(Koff[-1]) + return p_face + + +@field_operator +def _face_val_ppm_stencil_02( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + p_face_in: Field[[CellDim, KDim], float], + vert_idx: Field[[KDim], int32], + slev: int32, + elev: int32, + slevp1: int32, + elevp1: int32, +) -> Field[[CellDim, KDim], float]: + + vert_idx = broadcast(vert_idx, (CellDim, KDim)) + + p_face = where( + (vert_idx == slevp1) | (vert_idx == elev), + _face_val_ppm_stencil_02a(p_cc, p_cellhgt_mc_now), + p_face_in, + ) + + p_face = where((vert_idx == slev), _face_val_ppm_stencil_02b(p_cc), p_face) + + p_face = where((vert_idx == elevp1), _face_val_ppm_stencil_02c(p_cc), p_face) + + return p_face + + +@program +def face_val_ppm_stencil_02( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + p_face_in: Field[[CellDim, KDim], float], + vert_idx: Field[[KDim], int32], + slev: int32, + elev: int32, + slevp1: int32, + elevp1: int32, + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_02( + p_cc, + p_cellhgt_mc_now, + p_face_in, + vert_idx, + slev, + elev, + slevp1, + elevp1, + out=p_face, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py new file mode 100644 index 000000000..59a728184 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02a.py @@ -0,0 +1,43 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field + +from icon4py.model.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_02a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc * (1.0 - (p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1]))) + ( + p_cellhgt_mc_now / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) + ) * ((p_cellhgt_mc_now / p_cellhgt_mc_now(Koff[-1])) * p_cc + p_cc(Koff[-1])) + + return p_face + + +@program +def face_val_ppm_stencil_02a( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_02a( + p_cc, + p_cellhgt_mc_now, + out=p_face, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py new file mode 100644 index 000000000..bab061cc4 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02b.py @@ -0,0 +1,37 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field + +from icon4py.model.common.dimension import CellDim, KDim + + +@field_operator +def _face_val_ppm_stencil_02b( + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc + return p_face + + +@program +def face_val_ppm_stencil_02b( + p_cc: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_02b( + p_cc, + out=p_face, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py new file mode 100644 index 000000000..54821c903 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_02c.py @@ -0,0 +1,37 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field + +from icon4py.model.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_02c( + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + p_face = p_cc(Koff[-1]) + return p_face + + +@program +def face_val_ppm_stencil_02c( + p_cc: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_02c( + p_cc, + out=p_face, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py new file mode 100644 index 000000000..df6d76839 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/face_val_ppm_stencil_05.py @@ -0,0 +1,67 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field + +from icon4py.model.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _face_val_ppm_stencil_05( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + z_slope: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], float]: + + zgeo1 = p_cellhgt_mc_now(Koff[-1]) / (p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now) + zgeo2 = 1.0 / ( + p_cellhgt_mc_now(Koff[-2]) + + p_cellhgt_mc_now(Koff[-1]) + + p_cellhgt_mc_now + + p_cellhgt_mc_now(Koff[1]) + ) + zgeo3 = (p_cellhgt_mc_now(Koff[-2]) + p_cellhgt_mc_now(Koff[-1])) / ( + 2.0 * p_cellhgt_mc_now(Koff[-1]) + p_cellhgt_mc_now + ) + zgeo4 = (p_cellhgt_mc_now(Koff[1]) + p_cellhgt_mc_now) / ( + 2.0 * p_cellhgt_mc_now + p_cellhgt_mc_now(Koff[-1]) + ) + + p_face = ( + p_cc(Koff[-1]) + + zgeo1 * (p_cc - p_cc(Koff[-1])) + + zgeo2 + * ( + (2.0 * p_cellhgt_mc_now * zgeo1) * (zgeo3 - zgeo4) * (p_cc - p_cc(Koff[-1])) + - zgeo3 * p_cellhgt_mc_now(Koff[-1]) * z_slope + + zgeo4 * p_cellhgt_mc_now * z_slope(Koff[-1]) + ) + ) + + return p_face + + +@program +def face_val_ppm_stencil_05( + p_cc: Field[[CellDim, KDim], float], + p_cellhgt_mc_now: Field[[CellDim, KDim], float], + z_slope: Field[[CellDim, KDim], float], + p_face: Field[[CellDim, KDim], float], +): + _face_val_ppm_stencil_05( + p_cc, + p_cellhgt_mc_now, + z_slope, + out=p_face, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py new file mode 100644 index 000000000..7457dc3eb --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_01a.py @@ -0,0 +1,160 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, where + +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _hflux_ffsl_hybrid_stencil_01a( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + z_lsq_coeff_4: Field[[CellDim, KDim], float], + z_lsq_coeff_5: Field[[CellDim, KDim], float], + z_lsq_coeff_6: Field[[CellDim, KDim], float], + z_lsq_coeff_7: Field[[CellDim, KDim], float], + z_lsq_coeff_8: Field[[CellDim, KDim], float], + z_lsq_coeff_9: Field[[CellDim, KDim], float], + z_lsq_coeff_10: Field[[CellDim, KDim], float], + z_quad_vector_sum0_1: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_2: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_3: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_4: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_5: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_6: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_7: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_8: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_9: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_10: Field[[EdgeDim, KDim], float], + patch0_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], +) -> Field[[EdgeDim, KDim], float]: + + p_out_e_hybrid_1a = ( + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1(E2C[1]), + z_lsq_coeff_1(E2C[0]), + ) + * z_quad_vector_sum0_1 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2(E2C[1]), + z_lsq_coeff_2(E2C[0]), + ) + * z_quad_vector_sum0_2 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3(E2C[1]), + z_lsq_coeff_3(E2C[0]), + ) + * z_quad_vector_sum0_3 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4(E2C[1]), + z_lsq_coeff_4(E2C[0]), + ) + * z_quad_vector_sum0_4 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5(E2C[1]), + z_lsq_coeff_5(E2C[0]), + ) + * z_quad_vector_sum0_5 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6(E2C[1]), + z_lsq_coeff_6(E2C[0]), + ) + * z_quad_vector_sum0_6 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7(E2C[1]), + z_lsq_coeff_7(E2C[0]), + ) + * z_quad_vector_sum0_7 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8(E2C[1]), + z_lsq_coeff_8(E2C[0]), + ) + * z_quad_vector_sum0_8 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9(E2C[1]), + z_lsq_coeff_9(E2C[0]), + ) + * z_quad_vector_sum0_9 + + where( + patch0_cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10(E2C[1]), + z_lsq_coeff_10(E2C[0]), + ) + * z_quad_vector_sum0_10 + ) + + return p_out_e_hybrid_1a + + +@program +def hflux_ffsl_hybrid_stencil_01a( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + z_lsq_coeff_4: Field[[CellDim, KDim], float], + z_lsq_coeff_5: Field[[CellDim, KDim], float], + z_lsq_coeff_6: Field[[CellDim, KDim], float], + z_lsq_coeff_7: Field[[CellDim, KDim], float], + z_lsq_coeff_8: Field[[CellDim, KDim], float], + z_lsq_coeff_9: Field[[CellDim, KDim], float], + z_lsq_coeff_10: Field[[CellDim, KDim], float], + z_quad_vector_sum0_1: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_2: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_3: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_4: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_5: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_6: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_7: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_8: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_9: Field[[EdgeDim, KDim], float], + z_quad_vector_sum0_10: Field[[EdgeDim, KDim], float], + patch0_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], + p_out_e_hybrid_1a: Field[[EdgeDim, KDim], float], +): + _hflux_ffsl_hybrid_stencil_01a( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + z_lsq_coeff_4, + z_lsq_coeff_5, + z_lsq_coeff_6, + z_lsq_coeff_7, + z_lsq_coeff_8, + z_lsq_coeff_9, + z_lsq_coeff_10, + z_quad_vector_sum0_1, + z_quad_vector_sum0_2, + z_quad_vector_sum0_3, + z_quad_vector_sum0_4, + z_quad_vector_sum0_5, + z_quad_vector_sum0_6, + z_quad_vector_sum0_7, + z_quad_vector_sum0_8, + z_quad_vector_sum0_9, + z_quad_vector_sum0_10, + patch0_cell_rel_idx_dsl, + out=(p_out_e_hybrid_1a), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py new file mode 100644 index 000000000..fd8ef19b2 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflux_ffsl_hybrid_stencil_02.py @@ -0,0 +1,43 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field + +from icon4py.model.common.dimension import EdgeDim, KDim + + +@field_operator +def _hflux_ffsl_hybrid_stencil_02( + p_out_e_hybrid_2: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + z_dreg_area: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + + p_out_e_hybrid_2 = p_mass_flx_e * p_out_e_hybrid_2 / z_dreg_area + + return p_out_e_hybrid_2 + + +@program +def hflux_ffsl_hybrid_stencil_02( + p_out_e_hybrid_2: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + z_dreg_area: Field[[EdgeDim, KDim], float], +): + _hflux_ffsl_hybrid_stencil_02( + p_out_e_hybrid_2, + p_mass_flx_e, + z_dreg_area, + out=p_out_e_hybrid_2, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py new file mode 100644 index 000000000..f912fc1f4 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01a.py @@ -0,0 +1,50 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, abs + +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_01a( + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +) -> tuple[Field[[EdgeDim, KDim], float], Field[[EdgeDim, KDim], float]]: + + z_mflx_low = 0.5 * ( + p_mass_flx_e * (p_cc(E2C[0]) + p_cc(E2C[1])) + - abs(p_mass_flx_e) * (p_cc(E2C[1]) - p_cc(E2C[0])) + ) + + z_anti = p_mflx_tracer_h - z_mflx_low + + return (z_mflx_low, z_anti) + + +@program +def hflx_limiter_mo_stencil_01a( + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + z_anti: Field[[EdgeDim, KDim], float], +): + _hflx_limiter_mo_stencil_01a( + p_mflx_tracer_h, + p_mass_flx_e, + p_cc, + out=(z_mflx_low, z_anti), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py new file mode 100644 index 000000000..0c9c94324 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_01b.py @@ -0,0 +1,95 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum, neighbor_sum + +from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_01b( + geofac_div: Field[[CEDim], float], + p_rhodz_now: Field[[CellDim, KDim], float], + p_rhodz_new: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + z_anti: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + p_dtime: float, +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: + zero = broadcast(0.0, (CellDim, KDim)) + + z_mflx_anti_1 = p_dtime * geofac_div(C2CE[0]) / p_rhodz_new * z_anti(C2E[0]) + z_mflx_anti_2 = p_dtime * geofac_div(C2CE[1]) / p_rhodz_new * z_anti(C2E[1]) + z_mflx_anti_3 = p_dtime * geofac_div(C2CE[2]) / p_rhodz_new * z_anti(C2E[2]) + + z_mflx_anti_in = -1.0 * ( + minimum(zero, z_mflx_anti_1) + minimum(zero, z_mflx_anti_2) + minimum(zero, z_mflx_anti_3) + ) + + z_mflx_anti_out = ( + maximum(zero, z_mflx_anti_1) + maximum(zero, z_mflx_anti_2) + maximum(zero, z_mflx_anti_3) + ) + + z_fluxdiv_c = neighbor_sum(z_mflx_low(C2E) * geofac_div(C2CE), axis=C2EDim) + + z_tracer_new_low = (p_cc * p_rhodz_now - p_dtime * z_fluxdiv_c) / p_rhodz_new + z_tracer_max = maximum(p_cc, z_tracer_new_low) + z_tracer_min = minimum(p_cc, z_tracer_new_low) + + return ( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + ) + + +@program +def hflx_limiter_mo_stencil_01b( + geofac_div: Field[[CEDim], float], + p_rhodz_now: Field[[CellDim, KDim], float], + p_rhodz_new: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + z_anti: Field[[EdgeDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + p_dtime: float, + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], +): + _hflx_limiter_mo_stencil_01b( + geofac_div, + p_rhodz_now, + p_rhodz_new, + z_mflx_low, + z_anti, + p_cc, + p_dtime, + out=( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + ), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py new file mode 100644 index 000000000..6a6a9f47b --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_02.py @@ -0,0 +1,70 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, maximum, minimum, where + +from icon4py.model.common.dimension import CellDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_02( + refin_ctrl: Field[[CellDim], int32], + p_cc: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + lo_bound: int32, + hi_bound: int32, +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: + condition = (refin_ctrl == lo_bound) | (refin_ctrl == hi_bound) + z_tracer_new_out = where( + condition, + minimum(1.1 * p_cc, maximum(0.9 * p_cc, z_tracer_new_low)), + z_tracer_new_low, + ) + + z_tracer_max_out = where(condition, maximum(p_cc, z_tracer_new_out), z_tracer_max) + z_tracer_min_out = where(condition, minimum(p_cc, z_tracer_new_out), z_tracer_min) + + return (z_tracer_new_out, z_tracer_max_out, z_tracer_min_out) + + +@program +def hflx_limiter_mo_stencil_02( + refin_ctrl: Field[[CellDim], int32], + p_cc: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + lo_bound: int32, + hi_bound: int32, + z_tracer_new_low_out: Field[[CellDim, KDim], float], + z_tracer_max_out: Field[[CellDim, KDim], float], + z_tracer_min_out: Field[[CellDim, KDim], float], +): + _hflx_limiter_mo_stencil_02( + refin_ctrl, + p_cc, + z_tracer_new_low, + z_tracer_max, + z_tracer_min, + lo_bound, + hi_bound, + out=(z_tracer_new_low_out, z_tracer_max_out, z_tracer_min_out), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py new file mode 100644 index 000000000..a7184d105 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_03.py @@ -0,0 +1,114 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import max_over, maximum, min_over, minimum + +from icon4py.model.common.dimension import C2E2C, C2E2CDim, CellDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + z_max = beta_fct * maximum(max_over(z_tracer_max(C2E2C), axis=C2E2CDim), z_tracer_max) + z_min = r_beta_fct * minimum(min_over(z_tracer_min(C2E2C), axis=C2E2CDim), z_tracer_min) + return z_max, z_min + + +@program +def hflx_limiter_mo_stencil_03_min_max( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, + z_max: Field[[CellDim, KDim], float], + z_min: Field[[CellDim, KDim], float], +): + _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, z_tracer_min, beta_fct, r_beta_fct, out=(z_max, z_min) + ) + + +@field_operator +def _hflx_limiter_mo_stencil_03a( + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + z_max: Field[[CellDim, KDim], float], + z_min: Field[[CellDim, KDim], float], + dbl_eps: float, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + r_p = (z_max - z_tracer_new_low) / (z_mflx_anti_in + dbl_eps) + r_m = (z_tracer_new_low - z_min) / (z_mflx_anti_out + dbl_eps) + + return r_p, r_m + + +@field_operator +def _hflx_limiter_mo_stencil_03( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + dbl_eps: float, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + z_max, z_min = _hflx_limiter_mo_stencil_03_min_max( + z_tracer_max, z_tracer_min, beta_fct, r_beta_fct + ) + + r_p, r_m = _hflx_limiter_mo_stencil_03a( + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + z_max, + z_min, + dbl_eps, + ) + return r_p, r_m + + +@program +def hflx_limiter_mo_stencil_03( + z_tracer_max: Field[[CellDim, KDim], float], + z_tracer_min: Field[[CellDim, KDim], float], + beta_fct: float, + r_beta_fct: float, + z_mflx_anti_in: Field[[CellDim, KDim], float], + z_mflx_anti_out: Field[[CellDim, KDim], float], + z_tracer_new_low: Field[[CellDim, KDim], float], + dbl_eps: float, + r_p: Field[[CellDim, KDim], float], + r_m: Field[[CellDim, KDim], float], +): + + _hflx_limiter_mo_stencil_03( + z_tracer_max, + z_tracer_min, + beta_fct, + r_beta_fct, + z_mflx_anti_in, + z_mflx_anti_out, + z_tracer_new_low, + dbl_eps, + out=(r_p, r_m), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py new file mode 100644 index 000000000..76e196572 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py @@ -0,0 +1,44 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import minimum, where + +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_04( + z_anti: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + r_frac = where( + z_anti >= 0.0, + minimum(r_m(E2C[0]), r_p(E2C[1])), + minimum(r_m(E2C[1]), r_p(E2C[0])), + ) + return z_mflx_low + minimum(1.0, r_frac) * z_anti + + +@program +def hflx_limiter_mo_stencil_04( + z_anti: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], +): + _hflx_limiter_mo_stencil_04(z_anti, r_m, r_p, z_mflx_low, out=p_mflx_tracer_h) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py new file mode 100644 index 000000000..d645251f6 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_05.py @@ -0,0 +1,49 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import minimum, where + +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _hflx_limiter_mo_stencil_05( + z_anti: Field[[EdgeDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + + z_signum = where((z_anti > 0.0), 1.0, -1.0) + + r_frac = 0.5 * ( + (1.0 + z_signum) * minimum(r_m(E2C[0]), r_p(E2C[1])) + + (1.0 - z_signum) * minimum(r_m(E2C[1]), r_p(E2C[0])) + ) + + p_mflx_tracer_h = z_mflx_low + minimum(1.0, r_frac) * z_anti + + return p_mflx_tracer_h + + +@program +def hflx_limiter_mo_stencil_05( + z_anti: Field[[EdgeDim, KDim], float], + z_mflx_low: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + r_p: Field[[CellDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], +): + _hflx_limiter_mo_stencil_05(z_anti, z_mflx_low, r_m, r_p, out=p_mflx_tracer_h) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py new file mode 100644 index 000000000..39c253487 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_01.py @@ -0,0 +1,58 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, maximum, minimum + +from icon4py.model.common.dimension import C2CE, C2E, CEDim, CellDim, EdgeDim, KDim + + +@field_operator +def _hflx_limiter_pd_stencil_01( + geofac_div: Field[[CEDim], float], + p_cc: Field[[CellDim, KDim], float], + p_rhodz_now: Field[[CellDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + p_dtime: float, + dbl_eps: float, +) -> Field[[CellDim, KDim], float]: + zero = broadcast(0.0, (CellDim, KDim)) + + pm_0 = maximum(zero, p_mflx_tracer_h(C2E[0]) * geofac_div(C2CE[0]) * p_dtime) + pm_1 = maximum(zero, p_mflx_tracer_h(C2E[1]) * geofac_div(C2CE[1]) * p_dtime) + pm_2 = maximum(zero, p_mflx_tracer_h(C2E[2]) * geofac_div(C2CE[2]) * p_dtime) + p_m = pm_0 + pm_1 + pm_2 + r_m = minimum(broadcast(1.0, (CellDim, KDim)), (p_cc * p_rhodz_now) / (p_m + dbl_eps)) + + return r_m + + +@program +def hflx_limiter_pd_stencil_01( + geofac_div: Field[[CEDim], float], + p_cc: Field[[CellDim, KDim], float], + p_rhodz_now: Field[[CellDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + r_m: Field[[CellDim, KDim], float], + p_dtime: float, + dbl_eps: float, +): + _hflx_limiter_pd_stencil_01( + geofac_div, + p_cc, + p_rhodz_now, + p_mflx_tracer_h, + p_dtime, + dbl_eps, + out=r_m, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py new file mode 100644 index 000000000..9d1be4677 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py @@ -0,0 +1,52 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32, where + +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _hflx_limiter_pd_stencil_02( + refin_ctrl: Field[[EdgeDim], int32], + r_m: Field[[CellDim, KDim], float], + p_mflx_tracer_h_in: Field[[EdgeDim, KDim], float], + bound: int32, +) -> Field[[EdgeDim, KDim], float]: + p_mflx_tracer_h_out = where( + refin_ctrl == bound, + p_mflx_tracer_h_in, + where( + p_mflx_tracer_h_in >= 0.0, + p_mflx_tracer_h_in * r_m(E2C[0]), + p_mflx_tracer_h_in * r_m(E2C[1]), + ), + ) + return p_mflx_tracer_h_out + + +@program +def hflx_limiter_pd_stencil_02( + refin_ctrl: Field[[EdgeDim], int32], + r_m: Field[[CellDim, KDim], float], + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + bound: int32, +): + _hflx_limiter_pd_stencil_02( + refin_ctrl, + r_m, + p_mflx_tracer_h, + bound, + out=p_mflx_tracer_h, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py new file mode 100644 index 000000000..59ad061a9 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hor_adv_stencil_01.py @@ -0,0 +1,60 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.model.common.dimension import C2CE, C2E, C2EDim, CEDim, CellDim, EdgeDim, KDim + + +@field_operator +def _hor_adv_stencil_01( + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + deepatmo_divh: Field[[KDim], float], + tracer_now: Field[[CellDim, KDim], float], + rhodz_now: Field[[CellDim, KDim], float], + rhodz_new: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + tracer_new = ( + tracer_now * rhodz_now + - p_dtime + * deepatmo_divh + * neighbor_sum(p_mflx_tracer_h(C2E) * geofac_div(C2CE), axis=C2EDim) + ) / rhodz_new + + return tracer_new + + +@program +def hor_adv_stencil_01( + p_mflx_tracer_h: Field[[EdgeDim, KDim], float], + deepatmo_divh: Field[[KDim], float], + tracer_now: Field[[CellDim, KDim], float], + rhodz_now: Field[[CellDim, KDim], float], + rhodz_new: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + tracer_new: Field[[CellDim, KDim], float], + p_dtime: float, +): + _hor_adv_stencil_01( + p_mflx_tracer_h, + deepatmo_divh, + tracer_now, + rhodz_now, + rhodz_new, + geofac_div, + p_dtime, + out=tracer_new, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py similarity index 89% rename from model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py rename to model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py index cd8a2e9bf..572b0f8d6 100644 --- a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/mo_advection_traj_btraj_compute_o1_dsl.py +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/mo_advection_traj_btraj_compute_o1_dsl.py @@ -32,6 +32,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( dual_normal_cell_2: Field[[ECDim], float], p_dthalf: float, ) -> tuple[ + Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], int32], Field[[EdgeDim, KDim], float], @@ -40,6 +41,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( lvn_pos = where(p_vn > 0.0, True, False) p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1])) + p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1)) p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1])) z_ntdistv_bary_1 = -( @@ -66,7 +68,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl( + z_ntdistv_bary_2 * dual_normal_cell_2(E2EC[1]), ) - return p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2 + return p_cell_idx, p_cell_rel_idx_dsl, p_cell_blk, p_distv_bary_1, p_distv_bary_2 @program(grid_type=GridType.UNSTRUCTURED) @@ -82,6 +84,7 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2: Field[[ECDim], float], dual_normal_cell_2: Field[[ECDim], float], p_cell_idx: Field[[EdgeDim, KDim], int32], + p_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], p_cell_blk: Field[[EdgeDim, KDim], int32], p_distv_bary_1: Field[[EdgeDim, KDim], float], p_distv_bary_2: Field[[EdgeDim, KDim], float], @@ -99,5 +102,11 @@ def mo_advection_traj_btraj_compute_o1_dsl( primal_normal_cell_2, dual_normal_cell_2, p_dthalf, - out=(p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2), + out=( + p_cell_idx, + p_cell_rel_idx_dsl, + p_cell_blk, + p_distv_bary_1, + p_distv_bary_2, + ), ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py new file mode 100644 index 000000000..957ee74a0 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_list_stencil.py @@ -0,0 +1,435 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32, where + +from icon4py.model.common.dimension import EdgeDim, KDim + + +@field_operator +def _prep_gauss_quadrature_c_list_stencil( + famask_int: Field[[EdgeDim, KDim], int32], + p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_y: Field[[EdgeDim, KDim], float], + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_dreg_area_in: Field[[EdgeDim, KDim], float], +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], +]: + + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + z_eta_1_1 = 1.0 - eta_1 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 + z_eta_4_1 = 1.0 - eta_4 + z_eta_1_2 = 1.0 + eta_1 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 + z_eta_4_2 = 1.0 + eta_4 + z_eta_1_3 = 1.0 - zeta_1 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 + z_eta_4_3 = 1.0 - zeta_4 + z_eta_1_4 = 1.0 + zeta_1 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 + z_eta_4_4 = 1.0 + zeta_4 + + famask_bool = where(famask_int == int32(1), True, False) + p_coords_dreg_v_1_x = where(famask_bool, p_coords_dreg_v_1_x, 0.0) + p_coords_dreg_v_2_x = where(famask_bool, p_coords_dreg_v_2_x, 0.0) + p_coords_dreg_v_3_x = where(famask_bool, p_coords_dreg_v_3_x, 0.0) + p_coords_dreg_v_4_x = where(famask_bool, p_coords_dreg_v_4_x, 0.0) + p_coords_dreg_v_1_y = where(famask_bool, p_coords_dreg_v_1_y, 0.0) + p_coords_dreg_v_2_y = where(famask_bool, p_coords_dreg_v_2_y, 0.0) + p_coords_dreg_v_3_y = where(famask_bool, p_coords_dreg_v_3_y, 0.0) + p_coords_dreg_v_4_y = where(famask_bool, p_coords_dreg_v_4_y, 0.0) + + wgt_t_detjac_1 = where( + famask_bool, + dbl_eps + + z_wgt_1 + * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_2 = where( + famask_bool, + dbl_eps + + z_wgt_2 + * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_3 = where( + famask_bool, + dbl_eps + + z_wgt_3 + * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + wgt_t_detjac_4 = where( + famask_bool, + dbl_eps + + z_wgt_4 + * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ), + 0.0, + ) + + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + + p_dreg_area = p_dreg_area_in + p_quad_vector_sum_1 + + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area, + ) + + +@program +def prep_gauss_quadrature_c_list_stencil( + famask_int: Field[[EdgeDim, KDim], int32], + p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_y: Field[[EdgeDim, KDim], float], + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_dreg_area_in: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_1: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_2: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_3: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_4: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_5: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_6: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_7: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_8: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_9: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_10: Field[[EdgeDim, KDim], float], + p_dreg_area: Field[[EdgeDim, KDim], float], +): + _prep_gauss_quadrature_c_list_stencil( + famask_int, + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + p_dreg_area_in, + out=( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area, + ), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py new file mode 100644 index 000000000..d75ed51a3 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/prep_gauss_quadrature_c_stencil.py @@ -0,0 +1,396 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, abs, maximum, where + +from icon4py.model.common.dimension import EdgeDim, KDim + + +@field_operator +def _prep_gauss_quadrature_c_stencil( + p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_y: Field[[EdgeDim, KDim], float], + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], + Field[[EdgeDim, KDim], float], +]: + + z_wgt_1 = 0.0625 * wgt_zeta_1 * wgt_eta_1 + z_wgt_2 = 0.0625 * wgt_zeta_1 * wgt_eta_2 + z_wgt_3 = 0.0625 * wgt_zeta_2 * wgt_eta_1 + z_wgt_4 = 0.0625 * wgt_zeta_2 * wgt_eta_2 + + z_eta_1_1 = 1.0 - eta_1 + z_eta_2_1 = 1.0 - eta_2 + z_eta_3_1 = 1.0 - eta_3 + z_eta_4_1 = 1.0 - eta_4 + z_eta_1_2 = 1.0 + eta_1 + z_eta_2_2 = 1.0 + eta_2 + z_eta_3_2 = 1.0 + eta_3 + z_eta_4_2 = 1.0 + eta_4 + z_eta_1_3 = 1.0 - zeta_1 + z_eta_2_3 = 1.0 - zeta_2 + z_eta_3_3 = 1.0 - zeta_3 + z_eta_4_3 = 1.0 - zeta_4 + z_eta_1_4 = 1.0 + zeta_1 + z_eta_2_4 = 1.0 + zeta_2 + z_eta_3_4 = 1.0 + zeta_3 + z_eta_4_4 = 1.0 + zeta_4 + + wgt_t_detjac_1 = dbl_eps + z_wgt_1 * ( + ( + z_eta_1_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_1_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_1_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_1_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_1_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_1_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_1_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_2 = dbl_eps + z_wgt_2 * ( + ( + z_eta_2_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_2_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_2_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_2_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_2_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_2_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_2_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_3 = dbl_eps + z_wgt_3 * ( + ( + z_eta_3_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_3_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_3_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_3_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_3_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_3_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_3_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + wgt_t_detjac_4 = dbl_eps + z_wgt_4 * ( + ( + z_eta_4_1 * (p_coords_dreg_v_2_x - p_coords_dreg_v_1_x) + + z_eta_4_2 * (p_coords_dreg_v_3_x - p_coords_dreg_v_4_x) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_y - p_coords_dreg_v_1_y) + - z_eta_4_4 * (p_coords_dreg_v_2_y - p_coords_dreg_v_3_y) + ) + - ( + z_eta_4_1 * (p_coords_dreg_v_2_y - p_coords_dreg_v_1_y) + + z_eta_4_2 * (p_coords_dreg_v_3_y - p_coords_dreg_v_4_y) + ) + * ( + z_eta_4_3 * (p_coords_dreg_v_4_x - p_coords_dreg_v_1_x) + - z_eta_4_4 * (p_coords_dreg_v_2_x - p_coords_dreg_v_3_x) + ) + ) + + z_gauss_pts_1_x = ( + shape_func_1_1 * p_coords_dreg_v_1_x + + shape_func_2_1 * p_coords_dreg_v_2_x + + shape_func_3_1 * p_coords_dreg_v_3_x + + shape_func_4_1 * p_coords_dreg_v_4_x + ) + z_gauss_pts_1_y = ( + shape_func_1_1 * p_coords_dreg_v_1_y + + shape_func_2_1 * p_coords_dreg_v_2_y + + shape_func_3_1 * p_coords_dreg_v_3_y + + shape_func_4_1 * p_coords_dreg_v_4_y + ) + z_gauss_pts_2_x = ( + shape_func_1_2 * p_coords_dreg_v_1_x + + shape_func_2_2 * p_coords_dreg_v_2_x + + shape_func_3_2 * p_coords_dreg_v_3_x + + shape_func_4_2 * p_coords_dreg_v_4_x + ) + z_gauss_pts_2_y = ( + shape_func_1_2 * p_coords_dreg_v_1_y + + shape_func_2_2 * p_coords_dreg_v_2_y + + shape_func_3_2 * p_coords_dreg_v_3_y + + shape_func_4_2 * p_coords_dreg_v_4_y + ) + z_gauss_pts_3_x = ( + shape_func_1_3 * p_coords_dreg_v_1_x + + shape_func_2_3 * p_coords_dreg_v_2_x + + shape_func_3_3 * p_coords_dreg_v_3_x + + shape_func_4_3 * p_coords_dreg_v_4_x + ) + z_gauss_pts_3_y = ( + shape_func_1_3 * p_coords_dreg_v_1_y + + shape_func_2_3 * p_coords_dreg_v_2_y + + shape_func_3_3 * p_coords_dreg_v_3_y + + shape_func_4_3 * p_coords_dreg_v_4_y + ) + z_gauss_pts_4_x = ( + shape_func_1_4 * p_coords_dreg_v_1_x + + shape_func_2_4 * p_coords_dreg_v_2_x + + shape_func_3_4 * p_coords_dreg_v_3_x + + shape_func_4_4 * p_coords_dreg_v_4_x + ) + z_gauss_pts_4_y = ( + shape_func_1_4 * p_coords_dreg_v_1_y + + shape_func_2_4 * p_coords_dreg_v_2_y + + shape_func_3_4 * p_coords_dreg_v_3_y + + shape_func_4_4 * p_coords_dreg_v_4_y + ) + + p_quad_vector_sum_1 = wgt_t_detjac_1 + wgt_t_detjac_2 + wgt_t_detjac_3 + wgt_t_detjac_4 + p_quad_vector_sum_2 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x + ) + p_quad_vector_sum_3 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y + ) + p_quad_vector_sum_4 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_5 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_6 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_7 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_x + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_x + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_x + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_x + ) + p_quad_vector_sum_8 = ( + wgt_t_detjac_1 * z_gauss_pts_1_y * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_y * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_y * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_y * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + p_quad_vector_sum_9 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_x * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_x * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_x * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_x * z_gauss_pts_4_y + ) + p_quad_vector_sum_10 = ( + wgt_t_detjac_1 * z_gauss_pts_1_x * z_gauss_pts_1_y * z_gauss_pts_1_y + + wgt_t_detjac_2 * z_gauss_pts_2_x * z_gauss_pts_2_y * z_gauss_pts_2_y + + wgt_t_detjac_3 * z_gauss_pts_3_x * z_gauss_pts_3_y * z_gauss_pts_3_y + + wgt_t_detjac_4 * z_gauss_pts_4_x * z_gauss_pts_4_y * z_gauss_pts_4_y + ) + + z_area = p_quad_vector_sum_1 + p_dreg_area_out = where(z_area >= 0.0, maximum(eps, abs(z_area)), -maximum(eps, abs(z_area))) + + return ( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area_out, + ) + + +@program +def prep_gauss_quadrature_c_stencil( + p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_x: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_1_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_2_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_3_y: Field[[EdgeDim, KDim], float], + p_coords_dreg_v_4_y: Field[[EdgeDim, KDim], float], + shape_func_1_1: float, + shape_func_2_1: float, + shape_func_3_1: float, + shape_func_4_1: float, + shape_func_1_2: float, + shape_func_2_2: float, + shape_func_3_2: float, + shape_func_4_2: float, + shape_func_1_3: float, + shape_func_2_3: float, + shape_func_3_3: float, + shape_func_4_3: float, + shape_func_1_4: float, + shape_func_2_4: float, + shape_func_3_4: float, + shape_func_4_4: float, + zeta_1: float, + zeta_2: float, + zeta_3: float, + zeta_4: float, + eta_1: float, + eta_2: float, + eta_3: float, + eta_4: float, + wgt_zeta_1: float, + wgt_zeta_2: float, + wgt_eta_1: float, + wgt_eta_2: float, + dbl_eps: float, + eps: float, + p_quad_vector_sum_1: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_2: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_3: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_4: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_5: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_6: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_7: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_8: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_9: Field[[EdgeDim, KDim], float], + p_quad_vector_sum_10: Field[[EdgeDim, KDim], float], + p_dreg_area_out: Field[[EdgeDim, KDim], float], +): + _prep_gauss_quadrature_c_stencil( + p_coords_dreg_v_1_x, + p_coords_dreg_v_2_x, + p_coords_dreg_v_3_x, + p_coords_dreg_v_4_x, + p_coords_dreg_v_1_y, + p_coords_dreg_v_2_y, + p_coords_dreg_v_3_y, + p_coords_dreg_v_4_y, + shape_func_1_1, + shape_func_2_1, + shape_func_3_1, + shape_func_4_1, + shape_func_1_2, + shape_func_2_2, + shape_func_3_2, + shape_func_4_2, + shape_func_1_3, + shape_func_2_3, + shape_func_3_3, + shape_func_4_3, + shape_func_1_4, + shape_func_2_4, + shape_func_3_4, + shape_func_4_4, + zeta_1, + zeta_2, + zeta_3, + zeta_4, + eta_1, + eta_2, + eta_3, + eta_4, + wgt_zeta_1, + wgt_zeta_2, + wgt_eta_1, + wgt_eta_2, + dbl_eps, + eps, + out=( + p_quad_vector_sum_1, + p_quad_vector_sum_2, + p_quad_vector_sum_3, + p_quad_vector_sum_4, + p_quad_vector_sum_5, + p_quad_vector_sum_6, + p_quad_vector_sum_7, + p_quad_vector_sum_8, + p_quad_vector_sum_9, + p_quad_vector_sum_10, + p_dreg_area_out, + ), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/py.typed b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/py.typed new file mode 100644 index 000000000..e69de29bb diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py new file mode 100644 index 000000000..fa975ffcd --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/rbf_intp_edge_stencil_01.py @@ -0,0 +1,35 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.model.common.dimension import E2C2E, E2C2EDim, EdgeDim, KDim + + +@field_operator +def _rbf_intp_edge_stencil_01( + p_vn_in: Field[[EdgeDim, KDim], float], + ptr_coeff: Field[[EdgeDim, E2C2EDim], float], +) -> Field[[EdgeDim, KDim], float]: + p_vt_out = neighbor_sum(p_vn_in(E2C2E) * ptr_coeff, axis=E2C2EDim) + return p_vt_out + + +@program +def rbf_intp_edge_stencil_01( + p_vn_in: Field[[EdgeDim, KDim], float], + ptr_coeff: Field[[EdgeDim, E2C2EDim], float], + p_vt_out: Field[[EdgeDim, KDim], float], +): + _rbf_intp_edge_stencil_01(p_vn_in, ptr_coeff, out=p_vt_out) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py new file mode 100644 index 000000000..b16a6d96a --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_c_svd_stencil.py @@ -0,0 +1,244 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program + +from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim + + +@field_operator +def _recon_lsq_cell_c_svd_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECECDim], float], + lsq_pseudoinv_2: Field[[CECECDim], float], + lsq_pseudoinv_3: Field[[CECECDim], float], + lsq_pseudoinv_4: Field[[CECECDim], float], + lsq_pseudoinv_5: Field[[CECECDim], float], + lsq_pseudoinv_6: Field[[CECECDim], float], + lsq_pseudoinv_7: Field[[CECECDim], float], + lsq_pseudoinv_8: Field[[CECECDim], float], + lsq_pseudoinv_9: Field[[CECECDim], float], + lsq_moments_1: Field[[CellDim], float], + lsq_moments_2: Field[[CellDim], float], + lsq_moments_3: Field[[CellDim], float], + lsq_moments_4: Field[[CellDim], float], + lsq_moments_5: Field[[CellDim], float], + lsq_moments_6: Field[[CellDim], float], + lsq_moments_7: Field[[CellDim], float], + lsq_moments_8: Field[[CellDim], float], + lsq_moments_9: Field[[CellDim], float], +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: + + p_coeff_10_dsl = ( + lsq_pseudoinv_9(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_9(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_9_dsl = ( + lsq_pseudoinv_8(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_8(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_8_dsl = ( + lsq_pseudoinv_7(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_7(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_7_dsl = ( + lsq_pseudoinv_6(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_6(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_6_dsl = ( + lsq_pseudoinv_5(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_5(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_5_dsl = ( + lsq_pseudoinv_4(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_4(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_4_dsl = ( + lsq_pseudoinv_3(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_3(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_3_dsl = ( + lsq_pseudoinv_2(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_2(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + p_coeff_2_dsl = ( + lsq_pseudoinv_1(C2CECEC[0]) * (p_cc(C2E2C2E2C[0]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[1]) * (p_cc(C2E2C2E2C[1]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[2]) * (p_cc(C2E2C2E2C[2]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[3]) * (p_cc(C2E2C2E2C[3]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[4]) * (p_cc(C2E2C2E2C[4]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[5]) * (p_cc(C2E2C2E2C[5]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[6]) * (p_cc(C2E2C2E2C[6]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[7]) * (p_cc(C2E2C2E2C[7]) - p_cc) + + lsq_pseudoinv_1(C2CECEC[8]) * (p_cc(C2E2C2E2C[8]) - p_cc) + ) + + p_coeff_1_dsl = p_cc - ( + p_coeff_2_dsl * lsq_moments_1 + + p_coeff_3_dsl * lsq_moments_2 + + p_coeff_4_dsl * lsq_moments_3 + + p_coeff_5_dsl * lsq_moments_4 + + p_coeff_6_dsl * lsq_moments_5 + + p_coeff_7_dsl * lsq_moments_6 + + p_coeff_8_dsl * lsq_moments_7 + + p_coeff_9_dsl * lsq_moments_8 + + p_coeff_10_dsl * lsq_moments_9 + ) + + return ( + p_coeff_1_dsl, + p_coeff_2_dsl, + p_coeff_3_dsl, + p_coeff_4_dsl, + p_coeff_5_dsl, + p_coeff_6_dsl, + p_coeff_7_dsl, + p_coeff_8_dsl, + p_coeff_9_dsl, + p_coeff_10_dsl, + ) + + +@program +def recon_lsq_cell_c_svd_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECECDim], float], + lsq_pseudoinv_2: Field[[CECECDim], float], + lsq_pseudoinv_3: Field[[CECECDim], float], + lsq_pseudoinv_4: Field[[CECECDim], float], + lsq_pseudoinv_5: Field[[CECECDim], float], + lsq_pseudoinv_6: Field[[CECECDim], float], + lsq_pseudoinv_7: Field[[CECECDim], float], + lsq_pseudoinv_8: Field[[CECECDim], float], + lsq_pseudoinv_9: Field[[CECECDim], float], + lsq_moments_1: Field[[CellDim], float], + lsq_moments_2: Field[[CellDim], float], + lsq_moments_3: Field[[CellDim], float], + lsq_moments_4: Field[[CellDim], float], + lsq_moments_5: Field[[CellDim], float], + lsq_moments_6: Field[[CellDim], float], + lsq_moments_7: Field[[CellDim], float], + lsq_moments_8: Field[[CellDim], float], + lsq_moments_9: Field[[CellDim], float], + p_coeff_1_dsl: Field[[CellDim, KDim], float], + p_coeff_2_dsl: Field[[CellDim, KDim], float], + p_coeff_3_dsl: Field[[CellDim, KDim], float], + p_coeff_4_dsl: Field[[CellDim, KDim], float], + p_coeff_5_dsl: Field[[CellDim, KDim], float], + p_coeff_6_dsl: Field[[CellDim, KDim], float], + p_coeff_7_dsl: Field[[CellDim, KDim], float], + p_coeff_8_dsl: Field[[CellDim, KDim], float], + p_coeff_9_dsl: Field[[CellDim, KDim], float], + p_coeff_10_dsl: Field[[CellDim, KDim], float], +): + _recon_lsq_cell_c_svd_stencil( + p_cc, + lsq_pseudoinv_1, + lsq_pseudoinv_2, + lsq_pseudoinv_3, + lsq_pseudoinv_4, + lsq_pseudoinv_5, + lsq_pseudoinv_6, + lsq_pseudoinv_7, + lsq_pseudoinv_8, + lsq_pseudoinv_9, + lsq_moments_1, + lsq_moments_2, + lsq_moments_3, + lsq_moments_4, + lsq_moments_5, + lsq_moments_6, + lsq_moments_7, + lsq_moments_8, + lsq_moments_9, + out=( + p_coeff_1_dsl, + p_coeff_2_dsl, + p_coeff_3_dsl, + p_coeff_4_dsl, + p_coeff_5_dsl, + p_coeff_6_dsl, + p_coeff_7_dsl, + p_coeff_8_dsl, + p_coeff_9_dsl, + p_coeff_10_dsl, + ), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py new file mode 100644 index 000000000..5e8544755 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/recon_lsq_cell_l_svd_stencil.py @@ -0,0 +1,58 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program + +from icon4py.model.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim + + +@field_operator +def _recon_lsq_cell_l_svd_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECDim], float], + lsq_pseudoinv_2: Field[[CECDim], float], +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: + p_coeff_1_dsl = p_cc + p_coeff_2_dsl = ( + lsq_pseudoinv_1(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + + lsq_pseudoinv_1(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + + lsq_pseudoinv_1(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) + ) + p_coeff_3_dsl = ( + lsq_pseudoinv_2(C2CEC[0]) * (p_cc(C2E2C[0]) - p_cc) + + lsq_pseudoinv_2(C2CEC[1]) * (p_cc(C2E2C[1]) - p_cc) + + lsq_pseudoinv_2(C2CEC[2]) * (p_cc(C2E2C[2]) - p_cc) + ) + return p_coeff_1_dsl, p_coeff_2_dsl, p_coeff_3_dsl + + +@program +def recon_lsq_cell_l_svd_stencil( + p_cc: Field[[CellDim, KDim], float], + lsq_pseudoinv_1: Field[[CECDim], float], + lsq_pseudoinv_2: Field[[CECDim], float], + p_coeff_1_dsl: Field[[CellDim, KDim], float], + p_coeff_2_dsl: Field[[CellDim, KDim], float], + p_coeff_3_dsl: Field[[CellDim, KDim], float], +): + _recon_lsq_cell_l_svd_stencil( + p_cc, + lsq_pseudoinv_1, + lsq_pseudoinv_2, + out=(p_coeff_1_dsl, p_coeff_2_dsl, p_coeff_3_dsl), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py new file mode 100644 index 000000000..e0860c9fb --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c.py @@ -0,0 +1,28 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import broadcast + +from icon4py.model.common.dimension import CellDim + + +@field_operator +def _set_zero_c() -> Field[[CellDim], float]: + return broadcast(0.0, (CellDim,)) + + +@program +def set_zero_c(field: Field[[CellDim], float]): + _set_zero_c(out=field) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py new file mode 100644 index 000000000..f9db607ea --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/set_zero_c_k.py @@ -0,0 +1,28 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import broadcast + +from icon4py.model.common.dimension import CellDim, KDim + + +@field_operator +def _set_zero_c_k() -> Field[[CellDim, KDim], float]: + return broadcast(0.0, (CellDim, KDim)) + + +@program +def set_zero_c_k(field: Field[[CellDim, KDim], float]): + _set_zero_c_k(out=field) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py new file mode 100644 index 000000000..0ac10cc6a --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py @@ -0,0 +1,50 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program + +from icon4py.model.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _step_advection_stencil_01( + rhodz_ast: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + k_offset_up_low = p_dtime * ( + p_mflx_contra_v(Koff[1]) * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu + ) + return rhodz_ast + k_offset_up_low + + +@program +def step_advection_stencil_01( + rhodz_ast: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, + rhodz_ast2: Field[[CellDim, KDim], float], +): + _step_advection_stencil_01( + rhodz_ast, + p_mflx_contra_v, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + out=rhodz_ast2, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py new file mode 100644 index 000000000..af201f955 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py @@ -0,0 +1,50 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import maximum + +from icon4py.model.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _step_advection_stencil_02( + p_rhodz_new: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + return maximum(0.1 * p_rhodz_new, p_rhodz_new) - p_dtime * ( + p_mflx_contra_v(Koff[1]) * deepatmo_divzl - p_mflx_contra_v * deepatmo_divzu + ) + + +@program +def step_advection_stencil_02( + p_rhodz_new: Field[[CellDim, KDim], float], + p_mflx_contra_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + p_dtime: float, + rhodz_ast2: Field[[CellDim, KDim], float], +): + _step_advection_stencil_02( + p_rhodz_new, + p_mflx_contra_v, + deepatmo_divzl, + deepatmo_divzu, + p_dtime, + out=rhodz_ast2, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py new file mode 100644 index 000000000..b1e775c62 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py @@ -0,0 +1,37 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, maximum + +from icon4py.model.common.dimension import CellDim, KDim + + +@field_operator +def _step_advection_stencil_03( + p_tracer_now: Field[[CellDim, KDim], float], + p_grf_tend_tracer: Field[[CellDim, KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + p_tracer_new = maximum(0.0, p_tracer_now + p_dtime * p_grf_tend_tracer) + return p_tracer_new + + +@program +def step_advection_stencil_03( + p_tracer_now: Field[[CellDim, KDim], float], + p_grf_tend_tracer: Field[[CellDim, KDim], float], + p_tracer_new: Field[[CellDim, KDim], float], + p_dtime: float, +): + _step_advection_stencil_03(p_tracer_now, p_grf_tend_tracer, p_dtime, out=p_tracer_new) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py new file mode 100644 index 000000000..52e0c7940 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_04.py @@ -0,0 +1,37 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field + +from icon4py.model.common.dimension import CellDim, KDim + + +@field_operator +def _step_advection_stencil_04( + p_tracer_now: Field[[CellDim, KDim], float], + p_tracer_new: Field[[CellDim, KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + opt_ddt_tracer_adv = (p_tracer_new - p_tracer_now) / p_dtime + return opt_ddt_tracer_adv + + +@program +def step_advection_stencil_04( + p_tracer_now: Field[[CellDim, KDim], float], + p_tracer_new: Field[[CellDim, KDim], float], + opt_ddt_tracer_adv: Field[[CellDim, KDim], float], + p_dtime: float, +): + _step_advection_stencil_04(p_tracer_now, p_tracer_new, p_dtime, out=opt_ddt_tracer_adv) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py new file mode 100644 index 000000000..50a1be5a9 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura3_stencil_01.py @@ -0,0 +1,170 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, where + +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _upwind_hflux_miura3_stencil_01( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + z_lsq_coeff_4: Field[[CellDim, KDim], float], + z_lsq_coeff_5: Field[[CellDim, KDim], float], + z_lsq_coeff_6: Field[[CellDim, KDim], float], + z_lsq_coeff_7: Field[[CellDim, KDim], float], + z_lsq_coeff_8: Field[[CellDim, KDim], float], + z_lsq_coeff_9: Field[[CellDim, KDim], float], + z_lsq_coeff_10: Field[[CellDim, KDim], float], + z_quad_vector_sum_1: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_2: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_3: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_4: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_5: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_6: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_7: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_8: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_9: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_10: Field[[EdgeDim, KDim], float], + z_dreg_area: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], +) -> Field[[EdgeDim, KDim], float]: + + p_out_e_miura3 = ( + ( + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1(E2C[1]), + z_lsq_coeff_1(E2C[0]), + ) + * z_quad_vector_sum_1 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2(E2C[1]), + z_lsq_coeff_2(E2C[0]), + ) + * z_quad_vector_sum_2 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3(E2C[1]), + z_lsq_coeff_3(E2C[0]), + ) + * z_quad_vector_sum_3 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_4(E2C[1]), + z_lsq_coeff_4(E2C[0]), + ) + * z_quad_vector_sum_4 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_5(E2C[1]), + z_lsq_coeff_5(E2C[0]), + ) + * z_quad_vector_sum_5 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_6(E2C[1]), + z_lsq_coeff_6(E2C[0]), + ) + * z_quad_vector_sum_6 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_7(E2C[1]), + z_lsq_coeff_7(E2C[0]), + ) + * z_quad_vector_sum_7 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_8(E2C[1]), + z_lsq_coeff_8(E2C[0]), + ) + * z_quad_vector_sum_8 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_9(E2C[1]), + z_lsq_coeff_9(E2C[0]), + ) + * z_quad_vector_sum_9 + + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_10(E2C[1]), + z_lsq_coeff_10(E2C[0]), + ) + * z_quad_vector_sum_10 + ) + / z_dreg_area + * p_mass_flx_e + ) + + return p_out_e_miura3 + + +@program +def upwind_hflux_miura3_stencil_01( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + z_lsq_coeff_4: Field[[CellDim, KDim], float], + z_lsq_coeff_5: Field[[CellDim, KDim], float], + z_lsq_coeff_6: Field[[CellDim, KDim], float], + z_lsq_coeff_7: Field[[CellDim, KDim], float], + z_lsq_coeff_8: Field[[CellDim, KDim], float], + z_lsq_coeff_9: Field[[CellDim, KDim], float], + z_lsq_coeff_10: Field[[CellDim, KDim], float], + z_quad_vector_sum_1: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_2: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_3: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_4: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_5: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_6: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_7: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_8: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_9: Field[[EdgeDim, KDim], float], + z_quad_vector_sum_10: Field[[EdgeDim, KDim], float], + z_dreg_area: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], + p_out_e_miura3: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura3_stencil_01( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + z_lsq_coeff_4, + z_lsq_coeff_5, + z_lsq_coeff_6, + z_lsq_coeff_7, + z_lsq_coeff_8, + z_lsq_coeff_9, + z_lsq_coeff_10, + z_quad_vector_sum_1, + z_quad_vector_sum_2, + z_quad_vector_sum_3, + z_quad_vector_sum_4, + z_quad_vector_sum_5, + z_quad_vector_sum_6, + z_quad_vector_sum_7, + z_quad_vector_sum_8, + z_quad_vector_sum_9, + z_quad_vector_sum_10, + z_dreg_area, + p_mass_flx_e, + cell_rel_idx_dsl, + out=(p_out_e_miura3), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py new file mode 100644 index 000000000..eef9ba61e --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_01.py @@ -0,0 +1,75 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, where + +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _upwind_hflux_miura_cycl_stencil_01( + z_lsq_coeff_1_dsl: Field[[CellDim, KDim], float], + z_lsq_coeff_2_dsl: Field[[CellDim, KDim], float], + z_lsq_coeff_3_dsl: Field[[CellDim, KDim], float], + distv_bary_1: Field[[EdgeDim, KDim], float], + distv_bary_2: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], +) -> Field[[EdgeDim, KDim], float]: + + z_tracer_mflx_dsl = ( + where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_1_dsl(E2C[1]), + z_lsq_coeff_1_dsl(E2C[0]), + ) + + distv_bary_1 + * where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_2_dsl(E2C[1]), + z_lsq_coeff_2_dsl(E2C[0]), + ) + + distv_bary_2 + * where( + cell_rel_idx_dsl == int32(1), + z_lsq_coeff_3_dsl(E2C[1]), + z_lsq_coeff_3_dsl(E2C[0]), + ) + ) * p_mass_flx_e + + return z_tracer_mflx_dsl + + +@program +def upwind_hflux_miura_cycl_stencil_01( + z_lsq_coeff_1_dsl: Field[[CellDim, KDim], float], + z_lsq_coeff_2_dsl: Field[[CellDim, KDim], float], + z_lsq_coeff_3_dsl: Field[[CellDim, KDim], float], + distv_bary_1: Field[[EdgeDim, KDim], float], + distv_bary_2: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], + z_tracer_mflx_dsl: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura_cycl_stencil_01( + z_lsq_coeff_1_dsl, + z_lsq_coeff_2_dsl, + z_lsq_coeff_3_dsl, + distv_bary_1, + distv_bary_2, + p_mass_flx_e, + cell_rel_idx_dsl, + out=(z_tracer_mflx_dsl), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.py new file mode 100644 index 000000000..b0026495f --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_02.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 + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum + +from icon4py.model.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim + + +@field_operator +def _upwind_hflux_miura_cycl_stencil_02( + nsub: int32, + p_mass_flx_e: Field[[EdgeDim, KDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + z_rhofluxdiv_c: Field[[CellDim, KDim], float], + z_tracer_mflx: Field[[EdgeDim, KDim], float], + z_rho_now: Field[[CellDim, KDim], float], + z_tracer_now: Field[[CellDim, KDim], float], + z_dtsub: float, +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: + + z_rhofluxdiv_c_out = ( + neighbor_sum(p_mass_flx_e(C2E) * geofac_div, axis=C2EDim) + if nsub == int32(1) + else z_rhofluxdiv_c + ) + + z_fluxdiv_c_dsl = neighbor_sum(z_tracer_mflx(C2E) * geofac_div, axis=C2EDim) + + z_rho_new_dsl = z_rho_now - z_dtsub * z_rhofluxdiv_c_out + + z_tracer_new_dsl = (z_tracer_now * z_rho_now - z_dtsub * z_fluxdiv_c_dsl) / z_rho_new_dsl + + return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl) + + +@program +def upwind_hflux_miura_cycl_stencil_02( + nsub: int32, + p_mass_flx_e: Field[[EdgeDim, KDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + z_rhofluxdiv_c: Field[[CellDim, KDim], float], + z_tracer_mflx: Field[[EdgeDim, KDim], float], + z_rho_now: Field[[CellDim, KDim], float], + z_tracer_now: Field[[CellDim, KDim], float], + z_dtsub: float, + z_rhofluxdiv_c_out: Field[[CellDim, KDim], float], + z_fluxdiv_c_dsl: Field[[CellDim, KDim], float], + z_rho_new_dsl: Field[[CellDim, KDim], float], + z_tracer_new_dsl: Field[[CellDim, KDim], float], +): + _upwind_hflux_miura_cycl_stencil_02( + nsub, + p_mass_flx_e, + geofac_div, + z_rhofluxdiv_c, + z_tracer_mflx, + z_rho_now, + z_tracer_now, + z_dtsub, + out=(z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py new file mode 100644 index 000000000..53cdc339d --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py @@ -0,0 +1,36 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import float64 + +from icon4py.model.common.dimension import EdgeDim, KDim + + +@field_operator +def _upwind_hflux_miura_cycl_stencil_03a( + z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl) / float64(2) + return p_out_e + + +@program +def upwind_hflux_miura_cycl_stencil_03a( + z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], + p_out_e: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura_cycl_stencil_03a(z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl, out=(p_out_e)) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py new file mode 100644 index 000000000..4b5c73e40 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03b.py @@ -0,0 +1,39 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program + +from icon4py.model.common.dimension import EdgeDim, KDim + + +@field_operator +def _upwind_hflux_miura_cycl_stencil_03b( + z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_3_dsl: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + p_out_e = (z_tracer_mflx_1_dsl + z_tracer_mflx_2_dsl + z_tracer_mflx_3_dsl) / float(3) + return p_out_e + + +@program +def upwind_hflux_miura_cycl_stencil_03b( + z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float], + z_tracer_mflx_3_dsl: Field[[EdgeDim, KDim], float], + p_out_e: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura_cycl_stencil_03b( + z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl, z_tracer_mflx_3_dsl, out=(p_out_e) + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py new file mode 100644 index 000000000..3637ba94a --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_stencil_01.py @@ -0,0 +1,63 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import int32, where + +from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim + + +@field_operator +def _upwind_hflux_miura_stencil_01( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + distv_bary_1: Field[[EdgeDim, KDim], float], + distv_bary_2: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], +) -> Field[[EdgeDim, KDim], float]: + + p_out_e = ( + where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_1(E2C[1]), z_lsq_coeff_1(E2C[0])) + + distv_bary_1 + * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_2(E2C[1]), z_lsq_coeff_2(E2C[0])) + + distv_bary_2 + * where(cell_rel_idx_dsl == int32(1), z_lsq_coeff_3(E2C[1]), z_lsq_coeff_3(E2C[0])) + ) * p_mass_flx_e + + return p_out_e + + +@program +def upwind_hflux_miura_stencil_01( + z_lsq_coeff_1: Field[[CellDim, KDim], float], + z_lsq_coeff_2: Field[[CellDim, KDim], float], + z_lsq_coeff_3: Field[[CellDim, KDim], float], + distv_bary_1: Field[[EdgeDim, KDim], float], + distv_bary_2: Field[[EdgeDim, KDim], float], + p_mass_flx_e: Field[[EdgeDim, KDim], float], + cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32], + p_out_e: Field[[EdgeDim, KDim], float], +): + _upwind_hflux_miura_stencil_01( + z_lsq_coeff_1, + z_lsq_coeff_2, + z_lsq_coeff_3, + distv_bary_1, + distv_bary_2, + p_mass_flx_e, + cell_rel_idx_dsl, + out=(p_out_e), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py new file mode 100644 index 000000000..953b50770 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_vflux_ppm_stencil_01.py @@ -0,0 +1,40 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.common import Field +from gt4py.next.ffront.decorator import field_operator, program + +from icon4py.model.common.dimension import CellDim, KDim + + +@field_operator +def _upwind_vflux_ppm_stencil_01( + z_face_up: Field[[CellDim, KDim], float], + z_face_low: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + z_delta_q = 0.5 * (z_face_up - z_face_low) + z_a1 = p_cc - 0.5 * (z_face_up + z_face_low) + + return z_delta_q, z_a1 + + +@program +def upwind_vflux_ppm_stencil_01( + z_face_up: Field[[CellDim, KDim], float], + z_face_low: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + z_delta_q: Field[[CellDim, KDim], float], + z_a1: Field[[CellDim, KDim], float], +): + _upwind_vflux_ppm_stencil_01(z_face_up, z_face_low, p_cc, out=(z_delta_q, z_a1)) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.py new file mode 100644 index 000000000..3ded4ed63 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_01.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 + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, FieldOffset, abs, int32, where + +from icon4py.model.common.dimension import CellDim, KDim + + +Koff = FieldOffset("Koff", source=KDim, target=(KDim,)) + + +@field_operator +def _v_limit_prbl_sm_stencil_01( + p_face: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +) -> Field[[CellDim, KDim], int32]: + + z_delta = p_face - p_face(Koff[1]) + z_a6i = 6.0 * (p_cc - 0.5 * (p_face + p_face(Koff[1]))) + + l_limit = where(abs(z_delta) < -1.0 * z_a6i, int32(1), int32(0)) + + return l_limit + + +@program +def v_limit_prbl_sm_stencil_01( + p_face: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + l_limit: Field[[CellDim, KDim], int32], +): + _v_limit_prbl_sm_stencil_01( + p_face, + p_cc, + out=l_limit, + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py new file mode 100644 index 000000000..cba045918 --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/v_limit_prbl_sm_stencil_02.py @@ -0,0 +1,60 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, FieldOffset, int32, minimum, where + +from icon4py.model.common.dimension import CellDim, KDim + + +Koff = FieldOffset("Koff", source=KDim, target=(KDim,)) + + +@field_operator +def _v_limit_prbl_sm_stencil_02( + l_limit: Field[[CellDim, KDim], int32], + p_face: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + q_face_up, q_face_low = where( + l_limit != int32(0), + where( + (p_cc < minimum(p_face, p_face(Koff[1]))), + (p_cc, p_cc), + where( + p_face > p_face(Koff[1]), + (3.0 * p_cc - 2.0 * p_face(Koff[1]), p_face(Koff[1])), + (p_face, 3.0 * p_cc - 2.0 * p_face), + ), + ), + (p_face, p_face(Koff[1])), + ) + + return q_face_up, q_face_low + + +@program +def v_limit_prbl_sm_stencil_02( + l_limit: Field[[CellDim, KDim], int32], + p_face: Field[[CellDim, KDim], float], + p_cc: Field[[CellDim, KDim], float], + p_face_up: Field[[CellDim, KDim], float], + p_face_low: Field[[CellDim, KDim], float], +): + _v_limit_prbl_sm_stencil_02( + l_limit, + p_face, + p_cc, + out=(p_face_up, p_face_low), + ) diff --git a/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py new file mode 100644 index 000000000..78736914d --- /dev/null +++ b/model/atmosphere/advection/src/icon4py/model/atmosphere/advection/vert_adv_stencil_01.py @@ -0,0 +1,58 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field + +from icon4py.model.common.dimension import CellDim, KDim, Koff + + +@field_operator +def _vert_adv_stencil_01( + tracer_now: Field[[CellDim, KDim], float], + rhodz_now: Field[[CellDim, KDim], float], + p_mflx_tracer_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + rhodz_new: Field[[CellDim, KDim], float], + p_dtime: float, +) -> Field[[CellDim, KDim], float]: + tracer_new = ( + tracer_now * rhodz_now + + p_dtime * (p_mflx_tracer_v(Koff[1]) * deepatmo_divzl - p_mflx_tracer_v * deepatmo_divzu) + ) / rhodz_new + + return tracer_new + + +@program +def vert_adv_stencil_01( + tracer_now: Field[[CellDim, KDim], float], + rhodz_now: Field[[CellDim, KDim], float], + p_mflx_tracer_v: Field[[CellDim, KDim], float], + deepatmo_divzl: Field[[KDim], float], + deepatmo_divzu: Field[[KDim], float], + rhodz_new: Field[[CellDim, KDim], float], + tracer_new: Field[[CellDim, KDim], float], + p_dtime: float, +): + _vert_adv_stencil_01( + tracer_now, + rhodz_now, + p_mflx_tracer_v, + deepatmo_divzl, + deepatmo_divzu, + rhodz_new, + p_dtime, + out=tracer_new, + ) diff --git a/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py new file mode 100644 index 000000000..f23101a69 --- /dev/null +++ b/model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/compute_airmass.py @@ -0,0 +1,36 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field + +from icon4py.model.common.dimension import CellDim, KDim + + +@field_operator +def _compute_airmass( + rho_in: Field[[CellDim, KDim], float], + ddqz_z_full_in: Field[[CellDim, KDim], float], + deepatmo_t1mc_in: Field[[KDim], float], +) -> Field[[CellDim, KDim], float]: + return rho_in * ddqz_z_full_in * deepatmo_t1mc_in + + +@program +def compute_airmass( + rho_in: Field[[CellDim, KDim], float], + ddqz_z_full_in: Field[[CellDim, KDim], float], + deepatmo_t1mc_in: Field[[KDim], float], + airmass_out: Field[[CellDim, KDim], float], +): + _compute_airmass(rho_in, ddqz_z_full_in, deepatmo_t1mc_in, out=airmass_out) diff --git a/model/atmosphere/dycore/tests/test_compute_airmass.py b/model/atmosphere/dycore/tests/test_compute_airmass.py new file mode 100644 index 000000000..dab2dcb90 --- /dev/null +++ b/model/atmosphere/dycore/tests/test_compute_airmass.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 numpy as np + +from icon4py.model.atmosphere.dycore.compute_airmass import compute_airmass +from icon4py.model.common.dimension import CellDim, KDim +from icon4py.model.common.test_utils.helpers import random_field +from icon4py.model.common.test_utils.simple_mesh import SimpleMesh + + +def compute_airmass_numpy( + rho_in: np.array, ddqz_z_full_in: np.array, deepatmo_t1mc_in: np.array +) -> np.array: + return rho_in * ddqz_z_full_in * deepatmo_t1mc_in + + +def test_compute_airmass(): + mesh = SimpleMesh() + + rho_in = random_field(mesh, CellDim, KDim) + ddqz_z_full_in = random_field(mesh, CellDim, KDim) + deepatmo_t1mc_in = random_field(mesh, KDim) + airmass_out = random_field(mesh, CellDim, KDim) + + ref = compute_airmass_numpy( + np.asarray(rho_in), np.asarray(ddqz_z_full_in), np.asarray(deepatmo_t1mc_in) + ) + compute_airmass(rho_in, ddqz_z_full_in, deepatmo_t1mc_in, airmass_out, offset_provider={}) + assert np.allclose(airmass_out, ref) diff --git a/model/common/src/icon4py/model/common/dimension.py b/model/common/src/icon4py/model/common/dimension.py index fd6bb6111..49c74a68b 100644 --- a/model/common/src/icon4py/model/common/dimension.py +++ b/model/common/src/icon4py/model/common/dimension.py @@ -24,6 +24,8 @@ CECDim = Dimension("CEC") ECDim = Dimension("EC") ECVDim = Dimension("ECV") +CECDim = Dimension("CEC") +CECECDim = Dimension("CECEC") E2CDim = Dimension("E2C", DimensionKind.LOCAL) E2VDim = Dimension("E2V", DimensionKind.LOCAL) C2EDim = Dimension("C2E", DimensionKind.LOCAL) @@ -36,6 +38,7 @@ E2C2EODim = Dimension("E2C2EO", DimensionKind.LOCAL) E2C2EDim = Dimension("E2C2E", DimensionKind.LOCAL) C2E2CDim = Dimension("C2E2C", DimensionKind.LOCAL) +C2E2C2E2CDim = Dimension("C2E2C2E2C", DimensionKind.LOCAL) E2C = FieldOffset("E2C", source=CellDim, target=(EdgeDim, E2CDim)) C2E = FieldOffset("C2E", source=EdgeDim, target=(CellDim, C2EDim)) V2C = FieldOffset("V2C", source=CellDim, target=(VertexDim, V2CDim)) @@ -51,5 +54,9 @@ E2C2EO = FieldOffset("E2C2EO", source=EdgeDim, target=(EdgeDim, E2C2EODim)) E2C2E = FieldOffset("E2C2E", source=EdgeDim, target=(EdgeDim, E2C2EDim)) C2E2C = FieldOffset("C2E2C", source=CellDim, target=(CellDim, C2E2CDim)) +C2E2C2E2C = FieldOffset("C2E2C2E2C", source=CellDim, target=(CellDim, C2E2C2E2CDim)) +C2CEC = FieldOffset("C2CEC", source=CECDim, target=(CellDim, C2E2CDim)) +C2CECEC = FieldOffset("C2CECEC", source=CECECDim, target=(CellDim, C2E2C2E2CDim)) + Koff = FieldOffset("Koff", source=KDim, target=(KDim,)) KHalfOff = FieldOffset("KHalfOff", source=KHalfDim, target=(KHalfDim,)) 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 4715be0fb..763616cc7 100644 --- a/model/common/src/icon4py/model/common/test_utils/simple_mesh.py +++ b/model/common/src/icon4py/model/common/test_utils/simple_mesh.py @@ -17,6 +17,7 @@ from gt4py.next.iterator.embedded import NeighborTableOffsetProvider, StridedNeighborOffsetProvider from icon4py.model.common.dimension import ( + C2E2C2E2CDim, C2E2CDim, C2E2CODim, C2EDim, @@ -371,6 +372,29 @@ class SimpleMeshData: ] ) + c2e2c2e2c_table = np.asarray( + [ + [15, 4, 3, 12, 14, 1, 7, 6, 2], # 1c + [16, 5, 4, 12, 13, 2, 8, 7, 0], + [17, 3, 5, 13, 14, 0, 6, 8, 1], + [0, 6, 2, 17, 5, 9, 10, 15, 4], + [1, 7, 0, 15, 3, 16, 5, 10, 11], # 5c + [2, 8, 1, 4, 16, 17, 3, 9, 11], + [3, 10, 9, 2, 0, 7, 13, 8, 12], + [4, 11, 10, 0, 1, 8, 14, 6, 13], + [5, 9, 11, 1, 2, 3, 12, 7, 14], + [6, 12, 8, 5, 11, 3, 10, 16, 15], # 10c + [7, 13, 6, 3, 9, 4, 11, 16, 17], + [8, 14, 7, 4, 10, 5, 9, 15, 17], + [9, 16, 15, 8, 6, 1, 13, 0, 14], + [10, 17, 16, 6, 7, 2, 14, 1, 12], + [11, 15, 17, 7, 8, 2, 13, 0, 12], # 15c + [12, 0, 14, 11, 17, 9, 16, 3, 4], + [13, 1, 12, 9, 15, 10, 17, 4, 5], + [14, 2, 13, 10, 16, 5, 3, 11, 15], + ] + ) + class SimpleMesh: _DEFAULT_K_LEVEL = 10 @@ -388,6 +412,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.e2c2v = SimpleMeshData.e2c2v_table self.v2c = SimpleMeshData.v2c_table self.v2e = SimpleMeshData.v2e_table + self.c2e2c2e2c = SimpleMeshData.c2e2c2e2c_table self.n_e2c = self.e2c.shape[1] self.n_e2v = self.e2v.shape[1] self.n_c2e = self.c2e.shape[1] @@ -400,6 +425,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): self.n_c2v = self.c2v.shape[1] self.n_v2e = self.v2e.shape[1] self.n_cells = self.c2e.shape[0] + self.n_c2e2c2e2c = self.c2e2c2e2c.shape[1] self.n_edges = 27 self.n_vertices = 9 self.k_level = k_level @@ -420,6 +446,7 @@ def __init__(self, k_level: int = _DEFAULT_K_LEVEL): CEDim: self.n_cells * self.n_c2e, E2C2VDim: self.n_e2c2v, ECVDim: self.n_edges * self.n_e2c2v, + C2E2C2E2CDim: self.n_c2e2c2e2c, } def get_c2v_offset_provider(self) -> NeighborTableOffsetProvider: @@ -455,6 +482,9 @@ def get_e2v_offset_provider(self) -> NeighborTableOffsetProvider: def get_e2c2v_offset_provider(self) -> NeighborTableOffsetProvider: return NeighborTableOffsetProvider(self.e2c2v, EdgeDim, VertexDim, self.n_e2c2v) + def get_c2e2c2e2c_offset_provider(self) -> NeighborTableOffsetProvider: + return NeighborTableOffsetProvider(self.c2e2c2e2c, CellDim, CellDim, self.n_c2e2c2e2c) + def get_e2ecv_offset_provider(self): old_shape = self.e2c2v.shape e2ecv_table = np.arange(old_shape[0] * old_shape[1]).reshape(old_shape) @@ -479,6 +509,7 @@ def get_offset_provider(self): "E2C2V": self.get_e2c2v_offset_provider(), "C2CE": self.get_c2ce_offset_provider(), "Koff": KDim, + "C2E2C2E2C": self.get_c2e2c2e2c_offset_provider(), "E2ECV": StridedNeighborOffsetProvider(EdgeDim, ECVDim, self.n_e2c2v), "E2EC": StridedNeighborOffsetProvider(EdgeDim, ECDim, self.n_e2c), } diff --git a/model/requirements-dev.txt b/model/requirements-dev.txt index 206e27f11..d4c4c5727 100644 --- a/model/requirements-dev.txt +++ b/model/requirements-dev.txt @@ -1,5 +1,6 @@ -r ../base-requirements-dev.txt -e ./atmosphere/dycore -e ./atmosphere/diffusion +-e ./atmosphere/advection -e ./common -e ./driver diff --git a/model/requirements.txt b/model/requirements.txt index 519a20834..3330309e2 100644 --- a/model/requirements.txt +++ b/model/requirements.txt @@ -1,5 +1,6 @@ -r ../base-requirements.txt ./atmosphere/dycore ./atmosphere/diffusion +./atmosphere/advection ./common ./driver diff --git a/model/tox.ini b/model/tox.ini index 9ef46bd3c..d4a915637 100644 --- a/model/tox.ini +++ b/model/tox.ini @@ -14,8 +14,8 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/diffusion/src common/src driver/src - pytest -v -s -n auto --cov --cov-append --benchmark-disable + -pytest -v -s -n auto -m "not slow_tests" -cache-clear --cov --cov-reset --doctest-modules atmosphere/dycore/src atmosphere/diffusion/src atmosphere/advection/src common/src driver/src + pytest -v -s -n auto -m "not slow_tests" --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html -coverage html diff --git a/pytest.ini b/pytest.ini index 84a9d34b4..4c4723375 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,3 +1,5 @@ [pytest] norecursedirs = _external_src python_functions = test_* bench_* +markers = + slow_tests: mark test as slow. diff --git a/requirements-dev.txt b/requirements-dev.txt index ca807eac7..61f75ec63 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -2,6 +2,7 @@ # icon4py model -e ./model/atmosphere/dycore +-e ./model/atmosphere/advection -e ./model/common -e ./model/atmosphere/diffusion -e ./model/driver diff --git a/requirements.txt b/requirements.txt index 14c1cc762..656d6c7d2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,6 +3,7 @@ # icon4py model ./model/atmosphere/dycore ./model/atmosphere/diffusion +./model/atmosphere/advection ./model/common ./model/driver diff --git a/tools/pyproject.toml b/tools/pyproject.toml index 9df4270d6..06ee7aa9b 100644 --- a/tools/pyproject.toml +++ b/tools/pyproject.toml @@ -21,6 +21,7 @@ classifiers = [ dependencies = [ 'icon4py-atmosphere-diffusion', 'icon4py-atmosphere-dycore', + 'icon4py-atmosphere-advection', 'gt4py>=1.0.1', 'icon4py-common', 'tabulate>=0.8.9', diff --git a/tools/requirements-dev.txt b/tools/requirements-dev.txt index d4652678b..16afdad85 100644 --- a/tools/requirements-dev.txt +++ b/tools/requirements-dev.txt @@ -1,5 +1,6 @@ -r ../base-requirements-dev.txt -e ../model/atmosphere/dycore -e ../model/atmosphere/diffusion +-e ../model/atmosphere/advection -e ../model/common -e . diff --git a/tools/requirements.txt b/tools/requirements.txt index 27d1e077c..dfe3e8e4e 100644 --- a/tools/requirements.txt +++ b/tools/requirements.txt @@ -1,4 +1,5 @@ -r ../base-requirements.txt ../model/atmosphere/dycore +../model/atmosphere/advection ../model/common . diff --git a/tox.ini b/tox.ini index 4b8680a57..622d25dea 100644 --- a/tox.ini +++ b/tox.ini @@ -14,8 +14,8 @@ passenv = deps = -r {toxinidir}/requirements-dev.txt commands = - -pytest -v -s -n auto -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/diffusion/src model/driver/src model/common/src tools/src - pytest -v -s -n auto --cov --cov-append --benchmark-disable + -pytest -v -s -n auto -m "not slow_tests" -cache-clear --cov --cov-reset --doctest-modules model/atmosphere/dycore/src model/atmosphere/diffusion/src model/atmosphere/advection/src model/driver/src model/common/src tools/src + pytest -v -s -n auto -m "not slow_tests" --cov --cov-append --benchmark-disable commands_post = rm -rf tests/_reports/coverage_html -coverage html