From 9715db3df34e1250c51c7dbb7b2fdf8acca50d4e Mon Sep 17 00:00:00 2001 From: Will Sawyer Date: Thu, 11 Sep 2025 12:17:08 +0200 Subject: [PATCH 01/23] Corrected problems in write_fields, added code for full_muphys --- .../integration_tests/run_full_muphys.py | 60 +++++++++++++------ .../integration_tests/run_graupel_only.py | 16 ++--- 2 files changed, 52 insertions(+), 24 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 67f382a8b8..9639966f84 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -14,15 +14,22 @@ import argparse import sys +import time import gt4py.next as gtx import netCDF4 import numpy as np +try: + from netCDF4 import Dataset +except ImportError: + print("Netcdf not installed") + sys.exit() + from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations.graupel import ( graupel_run, ) -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.thermo import saturation_adjustment +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import saturation_adjustment from icon4py.model.common import dimension as dims, model_backends @@ -109,7 +116,7 @@ def calc_dz(ksize, z): def write_fields( output_filename, - ncells, + ncell, nlev, t, qv, @@ -125,7 +132,10 @@ def write_fields( pflx, pre_gsp, ): - ncfile = netCDF4.Dataset(output_filename, mode="w") + ncfile = Dataset(output_filename, mode="w") + ncells = ncfile.createDimension("ncells", ncell) + height = ncfile.createDimension("height", nlev) + height1 = ncfile.createDimension("height1", nlev+1) ta_var = ncfile.createVariable("ta", np.double, ("height", "ncells")) hus_var = ncfile.createVariable("hus", np.double, ("height", "ncells")) clw_var = ncfile.createVariable("clw", np.double, ("height", "ncells")) @@ -262,7 +272,15 @@ def write_fields( data.mask_out, ) -saturation_adjustment( +start_time = time.time() + +for _x in range(int(args.itime)+1): + + if _x == 1: # Only start timing second iteration + start_time = time.time() + + saturation_adjustment = saturation_adjustment.with_backend(model_backends.BACKENDS["gtfn_cpu"]) + saturation_adjustment( te=gtx.as_field( ( dims.CellDim, @@ -310,12 +328,12 @@ def write_fields( qce_out=qc_out, # Specific cloud water content mask_out=mask_out, # Mask of interest offset_provider={"Koff": dims.KDim}, -) + ) -ksize = data.dz.shape[0] -k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32)) -graupel_run = graupel_run.with_backend(model_backends.BACKENDS["gtfn_cpu"]) -graupel_run( + ksize = data.dz.shape[0] + k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32)) + graupel_run = graupel_run.with_backend(model_backends.BACKENDS["gtfn_cpu"]) + graupel_run( k=k, last_lev=ksize - 1, dz=gtx.as_field( @@ -404,15 +422,16 @@ def write_fields( pg=pg_out, pre=pre_out, offset_provider={"Koff": dims.KDim}, -) + ) -data.prr_gsp = np.transpose(pr_out[dims.KDim(ksize - 1)].asnumpy()) -data.prs_gsp = np.transpose(ps_out[dims.KDim(ksize - 1)].asnumpy()) -data.pri_gsp = np.transpose(pi_out[dims.KDim(ksize - 1)].asnumpy()) -data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) -data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) + data.prr_gsp = np.transpose(pr_out[dims.KDim(ksize - 1)].asnumpy()) + data.prs_gsp = np.transpose(ps_out[dims.KDim(ksize - 1)].asnumpy()) + data.pri_gsp = np.transpose(pi_out[dims.KDim(ksize - 1)].asnumpy()) + data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) + data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) -saturation_adjustment( + saturation_adjustment = saturation_adjustment.with_backend(model_backends.BACKENDS["gtfn_cpu"]) + saturation_adjustment( te=gtx.as_field( ( dims.CellDim, @@ -460,8 +479,14 @@ def write_fields( qce_out=qc_out, # Specific cloud water content mask_out=mask_out, # Mask of interest offset_provider={"Koff": dims.KDim}, -) + ) + if _x == int(args.itime): # End timer on last iteration + end_time = time.time() + + +elapsed_time = end_time - start_time +print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") write_fields( args.output_file, data.ncells, @@ -480,3 +505,4 @@ def write_fields( pflx=np.transpose(pflx_out.asnumpy()), pre_gsp=data.pre_gsp, ) + diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index 1bda067112..78c46f5725 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -111,7 +111,7 @@ def calc_dz(ksize, z): def write_fields( output_filename, - ncells, + ncell, nlev, t, qv, @@ -127,14 +127,17 @@ def write_fields( pflx, pre_gsp, ): - ncfile = Dataset(output_filename, mode="w") - ta_var = ncfile.createVariable("ta", np.double, ("height", "ncells")) + ncfile = Dataset(output_filename, mode="w") + ncells = ncfile.createDimension("ncells", ncell) + height = ncfile.createDimension("height", nlev) + height1 = ncfile.createDimension("height1", nlev+1) + ta_var = ncfile.createVariable("ta", np.double, ("height", "ncells")) hus_var = ncfile.createVariable("hus", np.double, ("height", "ncells")) clw_var = ncfile.createVariable("clw", np.double, ("height", "ncells")) cli_var = ncfile.createVariable("cli", np.double, ("height", "ncells")) - qr_var = ncfile.createVariable("qr", np.double, ("height", "ncells")) - qs_var = ncfile.createVariable("qs", np.double, ("height", "ncells")) - qg_var = ncfile.createVariable("qg", np.double, ("height", "ncells")) + qr_var = ncfile.createVariable("qr", np.double, ("height", "ncells")) + qs_var = ncfile.createVariable("qs", np.double, ("height", "ncells")) + qg_var = ncfile.createVariable("qg", np.double, ("height", "ncells")) pflx_var = ncfile.createVariable("pflx", np.double, ("height", "ncells")) prr_gsp_var = ncfile.createVariable("prr_gsp", np.double, ("height1", "ncells")) prs_gsp_var = ncfile.createVariable("prs_gsp", np.double, ("height1", "ncells")) @@ -157,7 +160,6 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() - args = get_args() set_lib_path(args.ldir) From 062e2da98b108ff91eddc66b2bb32aebc9858c13 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 11 Sep 2025 12:47:23 +0200 Subject: [PATCH 02/23] try cse fix for compile-time in saturation adjustment 'loop' --- .../test_saturation_adjustment.py | 11 ++++-- pyproject.toml | 2 +- uv.lock | 36 +++++++++---------- 3 files changed, 25 insertions(+), 24 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py index d83df25624..1e91248a57 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py @@ -5,6 +5,8 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from __future__ import annotations +from typing import TYPE_CHECKING import numpy as np import pytest @@ -14,15 +16,18 @@ from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.stencil_tests import StencilTest +if TYPE_CHECKING: + from icon4py.model.common.grid import base as base_grid -@pytest.mark.embedded_only + +# @pytest.mark.embedded_only class TestSaturationAdjustment(StencilTest): PROGRAM = saturation_adjustment OUTPUTS = ("te_out", "qve_out", "qce_out", "mask_out") @staticmethod def reference( - grid, + grid: base_grid.Grid, te: np.ndarray, qve: np.ndarray, qce: np.ndarray, @@ -39,7 +44,7 @@ def reference( ) @pytest.fixture - def input_data(self, grid): + def input_data(self, grid: base_grid.Grid) -> dict: return dict( te=data_alloc.constant_field( grid, 273.90911754406039, dims.CellDim, dims.KDim, dtype=wpfloat diff --git a/pyproject.toml b/pyproject.toml index 27cc353c95..5891a742d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -307,7 +307,7 @@ url = "https://test.pypi.org/simple/" [tool.uv.sources] dace = {git = "https://github.com/GridTools/dace", tag = "__gt4py-next-integration_2025_08_28"} # ghex = {git = "https://github.com/ghex-org/GHEX.git", branch = "master"} -# gt4py = {git = "https://github.com/GridTools/gt4py", branch = "main"} +gt4py = {git = "https://github.com/havogt/gt4py", branch = "cse_in_fieldop_fusion"} # gt4py = {index = "test.pypi"} icon4py-atmosphere-advection = {workspace = true} icon4py-atmosphere-diffusion = {workspace = true} diff --git a/uv.lock b/uv.lock index 34861433c3..1f0f2e14c0 100644 --- a/uv.lock +++ b/uv.lock @@ -1428,8 +1428,8 @@ wheels = [ [[package]] name = "gt4py" -version = "1.0.8" -source = { registry = "https://pypi.org/simple" } +version = "0.0.0+missing.version.info" +source = { git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion#5c90bcf21b379bf354e968217b4ad4edc2aab5be" } dependencies = [ { name = "attrs" }, { name = "black" }, @@ -1460,10 +1460,6 @@ dependencies = [ { name = "versioningit" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/45/83/0d546c3b8987ddb771a0a9780bc74688bdd2f702e925ea30da667506b72a/gt4py-1.0.8.tar.gz", hash = "sha256:1dd686836377dbcbd4d0c20ba4757fdb4ed605d3468c548a5dabb52aed5d3047", size = 723685, upload-time = "2025-09-04T07:22:23.997Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ae/6d/5b1119125f3a76f06160a745ff1e714440fa8dcc1123e637977b51f944d1/gt4py-1.0.8-py3-none-any.whl", hash = "sha256:4ad705ddbbaee8aed1a7a38748314ee6223b93ef705721c29bd9730eec3ac6c5", size = 924753, upload-time = "2025-09-04T07:22:22.23Z" }, -] [package.optional-dependencies] cuda11 = [ @@ -1805,7 +1801,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1822,7 +1818,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1839,7 +1835,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1856,7 +1852,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1873,7 +1869,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1943,10 +1939,10 @@ requires-dist = [ { name = "dace", marker = "extra == 'dace'", git = "https://github.com/GridTools/dace?tag=__gt4py-next-integration_2025_08_28" }, { name = "datashader", marker = "extra == 'io'", specifier = ">=0.16.1" }, { name = "ghex", marker = "extra == 'distributed'", specifier = ">=0.3.0" }, - { name = "gt4py", specifier = "==1.0.8" }, - { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'" }, - { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'" }, - { name = "gt4py", extras = ["next"], marker = "extra == 'dace'" }, + { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", extras = ["next"], marker = "extra == 'dace'", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, { name = "holoviews", marker = "extra == 'io'", specifier = ">=1.16.0" }, { name = "icon4py-common", extras = ["dace", "distributed", "io"], marker = "extra == 'all'", editable = "model/common" }, { name = "mpi4py", marker = "extra == 'distributed'", specifier = ">=3.1.5" }, @@ -1982,7 +1978,7 @@ dependencies = [ requires-dist = [ { name = "click", specifier = ">=8.0.1" }, { name = "devtools", specifier = ">=0.12" }, - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, { name = "icon4py-common", editable = "model/common" }, @@ -2010,7 +2006,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "filelock", specifier = ">=3.18.0" }, - { name = "gt4py", specifier = "==1.0.8" }, + { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, { name = "numpy", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, @@ -2056,9 +2052,9 @@ requires-dist = [ { name = "cupy-cuda11x", marker = "extra == 'cuda11'", specifier = ">=13.0" }, { name = "cupy-cuda12x", marker = "extra == 'cuda12'", specifier = ">=13.0" }, { name = "fprettify", specifier = ">=0.3.7" }, - { name = "gt4py", specifier = "==1.0.8" }, - { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'" }, - { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'" }, + { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, { name = "icon4py-atmosphere-advection", editable = "model/atmosphere/advection" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, From f4a73ec13c358cb860643be5a859baced6cf9607 Mon Sep 17 00:00:00 2001 From: William Sawyer Date: Tue, 16 Sep 2025 09:11:22 +0200 Subject: [PATCH 03/23] For benchmarking only, includes workarounds from Hannes; do not merge --- .../integration_tests/run_full_muphys.py | 46 ++++++- .../integration_tests/run_graupel_only.py | 130 +++++------------- 2 files changed, 76 insertions(+), 100 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 9639966f84..2d6b625bca 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -165,7 +165,7 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() - +backend = model_backends.BACKENDS["gtfn_gpu"] args = get_args() set_lib_path(args.ldir) @@ -179,6 +179,7 @@ def write_fields( dims.KDim, ), data.t_out, + allocator=backend, ) qv_out = gtx.as_field( ( @@ -186,6 +187,7 @@ def write_fields( dims.KDim, ), data.qv_out, + allocator=backend, ) qc_out = gtx.as_field( ( @@ -193,6 +195,7 @@ def write_fields( dims.KDim, ), data.qc_out, + allocator=backend, ) qr_out = gtx.as_field( ( @@ -200,6 +203,7 @@ def write_fields( dims.KDim, ), data.qr_out, + allocator=backend, ) qs_out = gtx.as_field( ( @@ -207,6 +211,7 @@ def write_fields( dims.KDim, ), data.qs_out, + allocator=backend, ) qi_out = gtx.as_field( ( @@ -214,6 +219,7 @@ def write_fields( dims.KDim, ), data.qi_out, + allocator=backend, ) qg_out = gtx.as_field( ( @@ -221,6 +227,7 @@ def write_fields( dims.KDim, ), data.qg_out, + allocator=backend, ) pflx_out = gtx.as_field( ( @@ -228,6 +235,7 @@ def write_fields( dims.KDim, ), data.pflx_out, + allocator=backend, ) pr_out = gtx.as_field( ( @@ -235,6 +243,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) ps_out = gtx.as_field( ( @@ -242,6 +251,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pi_out = gtx.as_field( ( @@ -249,6 +259,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pg_out = gtx.as_field( ( @@ -256,6 +267,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pre_out = gtx.as_field( ( @@ -263,6 +275,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) mask_out = gtx.as_field( ( @@ -270,6 +283,7 @@ def write_fields( dims.KDim, ), data.mask_out, + allocator=backend, ) start_time = time.time() @@ -279,7 +293,7 @@ def write_fields( if _x == 1: # Only start timing second iteration start_time = time.time() - saturation_adjustment = saturation_adjustment.with_backend(model_backends.BACKENDS["gtfn_cpu"]) + saturation_adjustment = saturation_adjustment.with_backend(backend) saturation_adjustment( te=gtx.as_field( ( @@ -287,6 +301,7 @@ def write_fields( dims.KDim, ), np.transpose(data.t[0, :, :]), + allocator=backend, ), qve=gtx.as_field( ( @@ -294,6 +309,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qv[0, :, :]), + allocator=backend, ), qce=gtx.as_field( ( @@ -301,6 +317,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qc[0, :, :]), + allocator=backend, ), qre=gtx.as_field( ( @@ -308,6 +325,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qr[0, :, :]), + allocator=backend, ), qti=gtx.as_field( ( @@ -315,6 +333,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qg[0, :, :] + data.qs[0, :, :] + data.qi[0, :, :]), + allocator=backend, ), # Total ice rho=gtx.as_field( ( @@ -322,6 +341,7 @@ def write_fields( dims.KDim, ), np.transpose(data.rho[0, :, :]), + allocator=backend, ), te_out=t_out, # Temperature qve_out=qv_out, # Specific humidity @@ -331,8 +351,8 @@ def write_fields( ) ksize = data.dz.shape[0] - k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32)) - graupel_run = graupel_run.with_backend(model_backends.BACKENDS["gtfn_cpu"]) + k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32), allocator=backend) + graupel_run = graupel_run.with_backend(backend) graupel_run( k=k, last_lev=ksize - 1, @@ -342,6 +362,7 @@ def write_fields( dims.KDim, ), np.transpose(data.dz[:, :]), + allocator=backend, ), te=gtx.as_field( ( @@ -349,6 +370,7 @@ def write_fields( dims.KDim, ), np.transpose(data.t[0, :, :]), + allocator=backend, ), p=gtx.as_field( ( @@ -356,6 +378,7 @@ def write_fields( dims.KDim, ), np.transpose(data.p[0, :, :]), + allocator=backend, ), rho=gtx.as_field( ( @@ -363,6 +386,7 @@ def write_fields( dims.KDim, ), np.transpose(data.rho[0, :, :]), + allocator=backend, ), qve=gtx.as_field( ( @@ -370,6 +394,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qv[0, :, :]), + allocator=backend, ), qce=gtx.as_field( ( @@ -377,6 +402,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qc[0, :, :]), + allocator=backend, ), qre=gtx.as_field( ( @@ -384,6 +410,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qr[0, :, :]), + allocator=backend, ), qse=gtx.as_field( ( @@ -391,6 +418,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qs[0, :, :]), + allocator=backend, ), qie=gtx.as_field( ( @@ -398,6 +426,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qi[0, :, :]), + allocator=backend, ), qge=gtx.as_field( ( @@ -405,6 +434,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qg[0, :, :]), + allocator=backend, ), dt=args.dt, qnc=args.qnc, @@ -430,7 +460,7 @@ def write_fields( data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) - saturation_adjustment = saturation_adjustment.with_backend(model_backends.BACKENDS["gtfn_cpu"]) + saturation_adjustment = saturation_adjustment.with_backend(backend) saturation_adjustment( te=gtx.as_field( ( @@ -438,6 +468,7 @@ def write_fields( dims.KDim, ), np.transpose(data.t[0, :, :]), + allocator=backend, ), qve=gtx.as_field( ( @@ -445,6 +476,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qv[0, :, :]), + allocator=backend, ), qce=gtx.as_field( ( @@ -452,6 +484,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qc[0, :, :]), + allocator=backend, ), qre=gtx.as_field( ( @@ -459,6 +492,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qr[0, :, :]), + allocator=backend, ), qti=gtx.as_field( ( @@ -466,6 +500,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qg[0, :, :] + data.qs[0, :, :] + data.qi[0, :, :]), + allocator=backend, ), # Total ice rho=gtx.as_field( ( @@ -473,6 +508,7 @@ def write_fields( dims.KDim, ), np.transpose(data.rho[0, :, :]), + allocator=backend, ), te_out=t_out, # Temperature qve_out=qv_out, # Specific humidity diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index 78c46f5725..a1d230f5ac 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -160,6 +160,7 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() +backend = model_backends.BACKENDS["dace_gpu"] args = get_args() set_lib_path(args.ldir) @@ -173,6 +174,7 @@ def write_fields( dims.KDim, ), data.t_out, + allocator=backend, ) qv_out = gtx.as_field( ( @@ -180,6 +182,7 @@ def write_fields( dims.KDim, ), data.qv_out, + allocator=backend, ) qc_out = gtx.as_field( ( @@ -187,6 +190,7 @@ def write_fields( dims.KDim, ), data.qc_out, + allocator=backend, ) qr_out = gtx.as_field( ( @@ -194,6 +198,7 @@ def write_fields( dims.KDim, ), data.qr_out, + allocator=backend, ) qs_out = gtx.as_field( ( @@ -201,6 +206,7 @@ def write_fields( dims.KDim, ), data.qs_out, + allocator=backend, ) qi_out = gtx.as_field( ( @@ -208,6 +214,7 @@ def write_fields( dims.KDim, ), data.qi_out, + allocator=backend, ) qg_out = gtx.as_field( ( @@ -215,6 +222,7 @@ def write_fields( dims.KDim, ), data.qg_out, + allocator=backend, ) pflx_out = gtx.as_field( ( @@ -222,6 +230,7 @@ def write_fields( dims.KDim, ), data.pflx_out, + allocator=backend, ) pr_out = gtx.as_field( ( @@ -229,6 +238,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) ps_out = gtx.as_field( ( @@ -236,6 +246,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pi_out = gtx.as_field( ( @@ -243,6 +254,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pg_out = gtx.as_field( ( @@ -250,6 +262,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pre_out = gtx.as_field( ( @@ -257,6 +270,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) mask_out = gtx.as_field( ( @@ -264,103 +278,17 @@ def write_fields( dims.KDim, ), data.mask_out, + allocator=backend, ) ksize = data.dz.shape[0] -k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32)) -graupel_run = graupel_run.with_backend(model_backends.BACKENDS["gtfn_cpu"]) -graupel_run( - k=k, - last_lev=ksize - 1, - dz=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.dz[:, :]), - ), - te=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.t[0, :, :]), - ), - p=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.p[0, :, :]), - ), - rho=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.rho[0, :, :]), - ), - qve=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qv[0, :, :]), - ), - qce=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qc[0, :, :]), - ), - qre=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qr[0, :, :]), - ), - qse=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qs[0, :, :]), - ), - qie=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qi[0, :, :]), - ), - qge=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qg[0, :, :]), - ), - dt=args.dt, - qnc=args.qnc, - t_out=t_out, - qv_out=qv_out, - qc_out=qc_out, - qr_out=qr_out, - qs_out=qs_out, - qi_out=qi_out, - qg_out=qg_out, - pflx=pflx_out, - pr=pr_out, - ps=ps_out, - pi=pi_out, - pg=pg_out, - pre=pre_out, - offset_provider={"Koff": dims.KDim}, -) -start_time = time.time() -for _x in range(int(args.itime)): +k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32), allocator=backend) + +for _x in range(int(args.itime)+1): + if _x == 1: # Only start timing second iteration + start_time = time.time() + + graupel_run = graupel_run.with_backend(backend) graupel_run( k=k, last_lev=ksize - 1, @@ -370,6 +298,7 @@ def write_fields( dims.KDim, ), np.transpose(data.dz[:, :]), + allocator=backend, ), te=gtx.as_field( ( @@ -377,6 +306,7 @@ def write_fields( dims.KDim, ), np.transpose(data.t[0, :, :]), + allocator=backend, ), p=gtx.as_field( ( @@ -384,6 +314,7 @@ def write_fields( dims.KDim, ), np.transpose(data.p[0, :, :]), + allocator=backend, ), rho=gtx.as_field( ( @@ -391,6 +322,7 @@ def write_fields( dims.KDim, ), np.transpose(data.rho[0, :, :]), + allocator=backend, ), qve=gtx.as_field( ( @@ -398,6 +330,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qv[0, :, :]), + allocator=backend, ), qce=gtx.as_field( ( @@ -405,6 +338,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qc[0, :, :]), + allocator=backend, ), qre=gtx.as_field( ( @@ -412,6 +346,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qr[0, :, :]), + allocator=backend, ), qse=gtx.as_field( ( @@ -419,6 +354,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qs[0, :, :]), + allocator=backend, ), qie=gtx.as_field( ( @@ -426,6 +362,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qi[0, :, :]), + allocator=backend, ), qge=gtx.as_field( ( @@ -433,6 +370,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qg[0, :, :]), + allocator=backend, ), dt=args.dt, qnc=args.qnc, @@ -451,7 +389,9 @@ def write_fields( pre=pre_out, offset_provider={"Koff": dims.KDim}, ) -end_time = time.time() + if _x == int(args.itime): # End timer on last iteration + end_time = time.time() + elapsed_time = end_time - start_time print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") From ff63fbba9e3496d1b7a043df1da3a8de8c57247c Mon Sep 17 00:00:00 2001 From: Will Sawyer Date: Tue, 16 Sep 2025 09:26:52 +0200 Subject: [PATCH 04/23] Revised integration tests to work with DaCe backends, credit to Edoardo --- .../integration_tests/run_full_muphys.py | 46 ++++++- .../integration_tests/run_graupel_only.py | 130 +++++------------- 2 files changed, 76 insertions(+), 100 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 9639966f84..2d6b625bca 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -165,7 +165,7 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() - +backend = model_backends.BACKENDS["gtfn_gpu"] args = get_args() set_lib_path(args.ldir) @@ -179,6 +179,7 @@ def write_fields( dims.KDim, ), data.t_out, + allocator=backend, ) qv_out = gtx.as_field( ( @@ -186,6 +187,7 @@ def write_fields( dims.KDim, ), data.qv_out, + allocator=backend, ) qc_out = gtx.as_field( ( @@ -193,6 +195,7 @@ def write_fields( dims.KDim, ), data.qc_out, + allocator=backend, ) qr_out = gtx.as_field( ( @@ -200,6 +203,7 @@ def write_fields( dims.KDim, ), data.qr_out, + allocator=backend, ) qs_out = gtx.as_field( ( @@ -207,6 +211,7 @@ def write_fields( dims.KDim, ), data.qs_out, + allocator=backend, ) qi_out = gtx.as_field( ( @@ -214,6 +219,7 @@ def write_fields( dims.KDim, ), data.qi_out, + allocator=backend, ) qg_out = gtx.as_field( ( @@ -221,6 +227,7 @@ def write_fields( dims.KDim, ), data.qg_out, + allocator=backend, ) pflx_out = gtx.as_field( ( @@ -228,6 +235,7 @@ def write_fields( dims.KDim, ), data.pflx_out, + allocator=backend, ) pr_out = gtx.as_field( ( @@ -235,6 +243,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) ps_out = gtx.as_field( ( @@ -242,6 +251,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pi_out = gtx.as_field( ( @@ -249,6 +259,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pg_out = gtx.as_field( ( @@ -256,6 +267,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pre_out = gtx.as_field( ( @@ -263,6 +275,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) mask_out = gtx.as_field( ( @@ -270,6 +283,7 @@ def write_fields( dims.KDim, ), data.mask_out, + allocator=backend, ) start_time = time.time() @@ -279,7 +293,7 @@ def write_fields( if _x == 1: # Only start timing second iteration start_time = time.time() - saturation_adjustment = saturation_adjustment.with_backend(model_backends.BACKENDS["gtfn_cpu"]) + saturation_adjustment = saturation_adjustment.with_backend(backend) saturation_adjustment( te=gtx.as_field( ( @@ -287,6 +301,7 @@ def write_fields( dims.KDim, ), np.transpose(data.t[0, :, :]), + allocator=backend, ), qve=gtx.as_field( ( @@ -294,6 +309,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qv[0, :, :]), + allocator=backend, ), qce=gtx.as_field( ( @@ -301,6 +317,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qc[0, :, :]), + allocator=backend, ), qre=gtx.as_field( ( @@ -308,6 +325,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qr[0, :, :]), + allocator=backend, ), qti=gtx.as_field( ( @@ -315,6 +333,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qg[0, :, :] + data.qs[0, :, :] + data.qi[0, :, :]), + allocator=backend, ), # Total ice rho=gtx.as_field( ( @@ -322,6 +341,7 @@ def write_fields( dims.KDim, ), np.transpose(data.rho[0, :, :]), + allocator=backend, ), te_out=t_out, # Temperature qve_out=qv_out, # Specific humidity @@ -331,8 +351,8 @@ def write_fields( ) ksize = data.dz.shape[0] - k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32)) - graupel_run = graupel_run.with_backend(model_backends.BACKENDS["gtfn_cpu"]) + k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32), allocator=backend) + graupel_run = graupel_run.with_backend(backend) graupel_run( k=k, last_lev=ksize - 1, @@ -342,6 +362,7 @@ def write_fields( dims.KDim, ), np.transpose(data.dz[:, :]), + allocator=backend, ), te=gtx.as_field( ( @@ -349,6 +370,7 @@ def write_fields( dims.KDim, ), np.transpose(data.t[0, :, :]), + allocator=backend, ), p=gtx.as_field( ( @@ -356,6 +378,7 @@ def write_fields( dims.KDim, ), np.transpose(data.p[0, :, :]), + allocator=backend, ), rho=gtx.as_field( ( @@ -363,6 +386,7 @@ def write_fields( dims.KDim, ), np.transpose(data.rho[0, :, :]), + allocator=backend, ), qve=gtx.as_field( ( @@ -370,6 +394,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qv[0, :, :]), + allocator=backend, ), qce=gtx.as_field( ( @@ -377,6 +402,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qc[0, :, :]), + allocator=backend, ), qre=gtx.as_field( ( @@ -384,6 +410,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qr[0, :, :]), + allocator=backend, ), qse=gtx.as_field( ( @@ -391,6 +418,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qs[0, :, :]), + allocator=backend, ), qie=gtx.as_field( ( @@ -398,6 +426,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qi[0, :, :]), + allocator=backend, ), qge=gtx.as_field( ( @@ -405,6 +434,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qg[0, :, :]), + allocator=backend, ), dt=args.dt, qnc=args.qnc, @@ -430,7 +460,7 @@ def write_fields( data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) - saturation_adjustment = saturation_adjustment.with_backend(model_backends.BACKENDS["gtfn_cpu"]) + saturation_adjustment = saturation_adjustment.with_backend(backend) saturation_adjustment( te=gtx.as_field( ( @@ -438,6 +468,7 @@ def write_fields( dims.KDim, ), np.transpose(data.t[0, :, :]), + allocator=backend, ), qve=gtx.as_field( ( @@ -445,6 +476,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qv[0, :, :]), + allocator=backend, ), qce=gtx.as_field( ( @@ -452,6 +484,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qc[0, :, :]), + allocator=backend, ), qre=gtx.as_field( ( @@ -459,6 +492,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qr[0, :, :]), + allocator=backend, ), qti=gtx.as_field( ( @@ -466,6 +500,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qg[0, :, :] + data.qs[0, :, :] + data.qi[0, :, :]), + allocator=backend, ), # Total ice rho=gtx.as_field( ( @@ -473,6 +508,7 @@ def write_fields( dims.KDim, ), np.transpose(data.rho[0, :, :]), + allocator=backend, ), te_out=t_out, # Temperature qve_out=qv_out, # Specific humidity diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index 78c46f5725..a1d230f5ac 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -160,6 +160,7 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() +backend = model_backends.BACKENDS["dace_gpu"] args = get_args() set_lib_path(args.ldir) @@ -173,6 +174,7 @@ def write_fields( dims.KDim, ), data.t_out, + allocator=backend, ) qv_out = gtx.as_field( ( @@ -180,6 +182,7 @@ def write_fields( dims.KDim, ), data.qv_out, + allocator=backend, ) qc_out = gtx.as_field( ( @@ -187,6 +190,7 @@ def write_fields( dims.KDim, ), data.qc_out, + allocator=backend, ) qr_out = gtx.as_field( ( @@ -194,6 +198,7 @@ def write_fields( dims.KDim, ), data.qr_out, + allocator=backend, ) qs_out = gtx.as_field( ( @@ -201,6 +206,7 @@ def write_fields( dims.KDim, ), data.qs_out, + allocator=backend, ) qi_out = gtx.as_field( ( @@ -208,6 +214,7 @@ def write_fields( dims.KDim, ), data.qi_out, + allocator=backend, ) qg_out = gtx.as_field( ( @@ -215,6 +222,7 @@ def write_fields( dims.KDim, ), data.qg_out, + allocator=backend, ) pflx_out = gtx.as_field( ( @@ -222,6 +230,7 @@ def write_fields( dims.KDim, ), data.pflx_out, + allocator=backend, ) pr_out = gtx.as_field( ( @@ -229,6 +238,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) ps_out = gtx.as_field( ( @@ -236,6 +246,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pi_out = gtx.as_field( ( @@ -243,6 +254,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pg_out = gtx.as_field( ( @@ -250,6 +262,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) pre_out = gtx.as_field( ( @@ -257,6 +270,7 @@ def write_fields( dims.KDim, ), np.zeros((data.ncells, data.nlev)), + allocator=backend, ) mask_out = gtx.as_field( ( @@ -264,103 +278,17 @@ def write_fields( dims.KDim, ), data.mask_out, + allocator=backend, ) ksize = data.dz.shape[0] -k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32)) -graupel_run = graupel_run.with_backend(model_backends.BACKENDS["gtfn_cpu"]) -graupel_run( - k=k, - last_lev=ksize - 1, - dz=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.dz[:, :]), - ), - te=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.t[0, :, :]), - ), - p=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.p[0, :, :]), - ), - rho=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.rho[0, :, :]), - ), - qve=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qv[0, :, :]), - ), - qce=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qc[0, :, :]), - ), - qre=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qr[0, :, :]), - ), - qse=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qs[0, :, :]), - ), - qie=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qi[0, :, :]), - ), - qge=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qg[0, :, :]), - ), - dt=args.dt, - qnc=args.qnc, - t_out=t_out, - qv_out=qv_out, - qc_out=qc_out, - qr_out=qr_out, - qs_out=qs_out, - qi_out=qi_out, - qg_out=qg_out, - pflx=pflx_out, - pr=pr_out, - ps=ps_out, - pi=pi_out, - pg=pg_out, - pre=pre_out, - offset_provider={"Koff": dims.KDim}, -) -start_time = time.time() -for _x in range(int(args.itime)): +k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32), allocator=backend) + +for _x in range(int(args.itime)+1): + if _x == 1: # Only start timing second iteration + start_time = time.time() + + graupel_run = graupel_run.with_backend(backend) graupel_run( k=k, last_lev=ksize - 1, @@ -370,6 +298,7 @@ def write_fields( dims.KDim, ), np.transpose(data.dz[:, :]), + allocator=backend, ), te=gtx.as_field( ( @@ -377,6 +306,7 @@ def write_fields( dims.KDim, ), np.transpose(data.t[0, :, :]), + allocator=backend, ), p=gtx.as_field( ( @@ -384,6 +314,7 @@ def write_fields( dims.KDim, ), np.transpose(data.p[0, :, :]), + allocator=backend, ), rho=gtx.as_field( ( @@ -391,6 +322,7 @@ def write_fields( dims.KDim, ), np.transpose(data.rho[0, :, :]), + allocator=backend, ), qve=gtx.as_field( ( @@ -398,6 +330,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qv[0, :, :]), + allocator=backend, ), qce=gtx.as_field( ( @@ -405,6 +338,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qc[0, :, :]), + allocator=backend, ), qre=gtx.as_field( ( @@ -412,6 +346,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qr[0, :, :]), + allocator=backend, ), qse=gtx.as_field( ( @@ -419,6 +354,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qs[0, :, :]), + allocator=backend, ), qie=gtx.as_field( ( @@ -426,6 +362,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qi[0, :, :]), + allocator=backend, ), qge=gtx.as_field( ( @@ -433,6 +370,7 @@ def write_fields( dims.KDim, ), np.transpose(data.qg[0, :, :]), + allocator=backend, ), dt=args.dt, qnc=args.qnc, @@ -451,7 +389,9 @@ def write_fields( pre=pre_out, offset_provider={"Koff": dims.KDim}, ) -end_time = time.time() + if _x == int(args.itime): # End timer on last iteration + end_time = time.time() + elapsed_time = end_time - start_time print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") From 77537ff76b6693fb583d53b75fc4d22476dc06b7 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 16 Sep 2025 12:18:05 +0200 Subject: [PATCH 05/23] rewrite run_full_muphys --- .../integration_tests/run_full_muphys.py | 435 ++++++++---------- 1 file changed, 196 insertions(+), 239 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 2d6b625bca..dd2f23a6ae 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -20,16 +20,19 @@ import netCDF4 import numpy as np +from icon4py.model.common.model_options import setup_program + + try: from netCDF4 import Dataset except ImportError: print("Netcdf not installed") sys.exit() +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import saturation_adjustment from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations.graupel import ( graupel_run, ) -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import saturation_adjustment from icon4py.model.common import dimension as dims, model_backends @@ -46,6 +49,7 @@ def get_args(): help="output filename", default="output.nc", ) + parser.add_argument("backend", help="gt4py backend", default="gtfn_cpu") parser.add_argument("input_file", help="input data file") parser.add_argument("itime", help="time-index", nargs="?", default=0) parser.add_argument("dt", help="timestep", nargs="?", default=30.0) @@ -114,6 +118,194 @@ def calc_dz(ksize, z): return dz +def run_program(args, backend, data): + ksize = data.dz.shape[0] + + dz = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.dz[:, :]), + allocator=backend, + ) + te = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.t[0, :, :]), + allocator=backend, + ) + p = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.p[0, :, :]), + allocator=backend, + ) + qse = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qs[0, :, :]), + allocator=backend, + ) + qie = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qi[0, :, :]), + allocator=backend, + ) + qge = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qg[0, :, :]), + allocator=backend, + ) + qve = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qv[0, :, :]), + allocator=backend, + ) + qce = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qc[0, :, :]), + allocator=backend, + ) + qre = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qr[0, :, :]), + allocator=backend, + ) + qti = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qg[0, :, :] + data.qs[0, :, :] + data.qi[0, :, :]), + allocator=backend, + ) # Total ice + rho = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.rho[0, :, :]), + allocator=backend, + ) + + k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32), allocator=backend) + + saturation_adjustment_program = setup_program( + backend=backend, + program=saturation_adjustment, + constant_args={}, + horizontal_sizes={}, + vertical_sizes={}, + offset_provider={"Koff": dims.KDim}, + ) + + graupel_run_program = setup_program( + backend=backend, + program=graupel_run, + constant_args={}, + horizontal_sizes={}, + vertical_sizes={ + "last_lev": ksize - 1, + }, + offset_provider={"Koff": dims.KDim}, + ) + + start_time = time.time() + + for _x in range(int(args.itime) + 1): + if _x == 1: # Only start timing second iteration + start_time = time.time() + + saturation_adjustment_program( + te=te, + qve=qve, + qce=qce, + qre=qre, + qti=qti, # Total ice + rho=rho, + te_out=t_out, # Temperature + qve_out=qv_out, # Specific humidity + qce_out=qc_out, # Specific cloud water content + mask_out=mask_out, # Mask of interest + ) + + graupel_run_program( + k=k, + dz=dz, + te=te, + p=p, + rho=rho, + qve=qve, + qce=qce, + qre=qre, + qse=qse, + qie=qie, + qge=qge, + dt=args.dt, + qnc=args.qnc, + t_out=t_out, + qv_out=qv_out, + qc_out=qc_out, + qr_out=qr_out, + qs_out=qs_out, + qi_out=qi_out, + qg_out=qg_out, + pflx=pflx_out, + pr=pr_out, + ps=ps_out, + pi=pi_out, + pg=pg_out, + pre=pre_out, + ) + + data.prr_gsp = np.transpose(pr_out[dims.KDim(ksize - 1)].asnumpy()) + data.prs_gsp = np.transpose(ps_out[dims.KDim(ksize - 1)].asnumpy()) + data.pri_gsp = np.transpose(pi_out[dims.KDim(ksize - 1)].asnumpy()) + data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) + data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) + + saturation_adjustment_program( + te=te, + qve=qve, + qce=qce, + qre=qre, + qti=qti, # Total ice + rho=rho, + te_out=t_out, # Temperature + qve_out=qv_out, # Specific humidity + qce_out=qc_out, # Specific cloud water content + mask_out=mask_out, # Mask of interest + ) + + if _x == int(args.itime): # End timer on last iteration + end_time = time.time() + + elapsed_time = end_time - start_time + print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") + + def write_fields( output_filename, ncell, @@ -165,8 +357,9 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() -backend = model_backends.BACKENDS["gtfn_gpu"] + args = get_args() +backend = model_backends.BACKENDS[args.backend] set_lib_path(args.ldir) sys.setrecursionlimit(10**4) @@ -286,243 +479,8 @@ def write_fields( allocator=backend, ) -start_time = time.time() - -for _x in range(int(args.itime)+1): - - if _x == 1: # Only start timing second iteration - start_time = time.time() +run_program(args, backend, data) - saturation_adjustment = saturation_adjustment.with_backend(backend) - saturation_adjustment( - te=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.t[0, :, :]), - allocator=backend, - ), - qve=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qv[0, :, :]), - allocator=backend, - ), - qce=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qc[0, :, :]), - allocator=backend, - ), - qre=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qr[0, :, :]), - allocator=backend, - ), - qti=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qg[0, :, :] + data.qs[0, :, :] + data.qi[0, :, :]), - allocator=backend, - ), # Total ice - rho=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.rho[0, :, :]), - allocator=backend, - ), - te_out=t_out, # Temperature - qve_out=qv_out, # Specific humidity - qce_out=qc_out, # Specific cloud water content - mask_out=mask_out, # Mask of interest - offset_provider={"Koff": dims.KDim}, - ) - - ksize = data.dz.shape[0] - k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32), allocator=backend) - graupel_run = graupel_run.with_backend(backend) - graupel_run( - k=k, - last_lev=ksize - 1, - dz=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.dz[:, :]), - allocator=backend, - ), - te=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.t[0, :, :]), - allocator=backend, - ), - p=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.p[0, :, :]), - allocator=backend, - ), - rho=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.rho[0, :, :]), - allocator=backend, - ), - qve=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qv[0, :, :]), - allocator=backend, - ), - qce=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qc[0, :, :]), - allocator=backend, - ), - qre=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qr[0, :, :]), - allocator=backend, - ), - qse=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qs[0, :, :]), - allocator=backend, - ), - qie=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qi[0, :, :]), - allocator=backend, - ), - qge=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qg[0, :, :]), - allocator=backend, - ), - dt=args.dt, - qnc=args.qnc, - t_out=t_out, - qv_out=qv_out, - qc_out=qc_out, - qr_out=qr_out, - qs_out=qs_out, - qi_out=qi_out, - qg_out=qg_out, - pflx=pflx_out, - pr=pr_out, - ps=ps_out, - pi=pi_out, - pg=pg_out, - pre=pre_out, - offset_provider={"Koff": dims.KDim}, - ) - - data.prr_gsp = np.transpose(pr_out[dims.KDim(ksize - 1)].asnumpy()) - data.prs_gsp = np.transpose(ps_out[dims.KDim(ksize - 1)].asnumpy()) - data.pri_gsp = np.transpose(pi_out[dims.KDim(ksize - 1)].asnumpy()) - data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) - data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) - - saturation_adjustment = saturation_adjustment.with_backend(backend) - saturation_adjustment( - te=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.t[0, :, :]), - allocator=backend, - ), - qve=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qv[0, :, :]), - allocator=backend, - ), - qce=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qc[0, :, :]), - allocator=backend, - ), - qre=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qr[0, :, :]), - allocator=backend, - ), - qti=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qg[0, :, :] + data.qs[0, :, :] + data.qi[0, :, :]), - allocator=backend, - ), # Total ice - rho=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.rho[0, :, :]), - allocator=backend, - ), - te_out=t_out, # Temperature - qve_out=qv_out, # Specific humidity - qce_out=qc_out, # Specific cloud water content - mask_out=mask_out, # Mask of interest - offset_provider={"Koff": dims.KDim}, - ) - - if _x == int(args.itime): # End timer on last iteration - end_time = time.time() - - -elapsed_time = end_time - start_time -print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") write_fields( args.output_file, data.ncells, @@ -541,4 +499,3 @@ def write_fields( pflx=np.transpose(pflx_out.asnumpy()), pre_gsp=data.pre_gsp, ) - From f8c22e5b1b89b0cff3bb710490891145e1da6b1c Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 16 Sep 2025 12:21:02 +0200 Subject: [PATCH 06/23] fix backend arg --- .../tests/muphys/integration_tests/run_full_muphys.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index dd2f23a6ae..76bf910897 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -49,7 +49,13 @@ def get_args(): help="output filename", default="output.nc", ) - parser.add_argument("backend", help="gt4py backend", default="gtfn_cpu") + parser.add_argument( + "-b", + metavar="backend", + dest="backend", + help="gt4py backend", + default="gtfn_cpu", + ) parser.add_argument("input_file", help="input data file") parser.add_argument("itime", help="time-index", nargs="?", default=0) parser.add_argument("dt", help="timestep", nargs="?", default=30.0) From 546bb753ba3f523e37860a26398539619b87c5fb Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 16 Sep 2025 12:32:44 +0200 Subject: [PATCH 07/23] remove k_field --- .../muphys/implementations/graupel.py | 6 +----- .../tests/muphys/integration_tests/run_full_muphys.py | 3 --- .../tests/muphys/integration_tests/run_graupel_only.py | 8 +++----- 3 files changed, 4 insertions(+), 13 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py index 6e9b3e8d1d..2db883f645 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py @@ -365,7 +365,6 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] @gtx.field_operator def _precipitation_effects( - k: fa.KField[gtx.int32], last_lev: gtx.int32, kmin_r: fa.CellKField[bool], # rain minimum level kmin_i: fa.CellKField[bool], # ice minimum level @@ -437,7 +436,6 @@ def _precipitation_effects( @gtx.field_operator def _graupel_run( - k: fa.KField[gtx.int32], last_lev: gtx.int32, dz: fa.CellKField[ta.wpfloat], te: fa.CellKField[ta.wpfloat], # Temperature @@ -473,7 +471,7 @@ def _graupel_run( te, p, rho, qve, qce, qre, qse, qie, qge, mask, is_sig_present, dt, qnc ) qr, qs, qi, qg, t, pflx, pr, ps, pi, pg, pre = _precipitation_effects( - k, last_lev, kmin_r, kmin_i, kmin_s, kmin_g, qv, qc, qr, qs, qi, qg, t, rho, dz, dt + last_lev, kmin_r, kmin_i, kmin_s, kmin_g, qv, qc, qr, qs, qi, qg, t, rho, dz, dt ) return t, qv, qc, qr, qs, qi, qg, pflx, pr, ps, pi, pg, pre @@ -481,7 +479,6 @@ def _graupel_run( @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def graupel_run( - k: fa.KField[gtx.int32], last_lev: gtx.int32, dz: fa.CellKField[ta.wpfloat], te: fa.CellKField[ta.wpfloat], # Temperature @@ -510,7 +507,6 @@ def graupel_run( pre: fa.CellKField[ta.wpfloat], # Precipitation of graupel ): _graupel_run( - k, last_lev, dz, te, diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 76bf910897..95af1590f1 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -216,8 +216,6 @@ def run_program(args, backend, data): allocator=backend, ) - k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32), allocator=backend) - saturation_adjustment_program = setup_program( backend=backend, program=saturation_adjustment, @@ -258,7 +256,6 @@ def run_program(args, backend, data): ) graupel_run_program( - k=k, dz=dz, te=te, p=p, diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index a1d230f5ac..fdd2009f2e 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -282,15 +282,13 @@ def write_fields( ) ksize = data.dz.shape[0] -k = gtx.as_field((dims.KDim,), np.arange(0, ksize, dtype=np.int32), allocator=backend) -for _x in range(int(args.itime)+1): - if _x == 1: # Only start timing second iteration +for _x in range(int(args.itime) + 1): + if _x == 1: # Only start timing second iteration start_time = time.time() graupel_run = graupel_run.with_backend(backend) graupel_run( - k=k, last_lev=ksize - 1, dz=gtx.as_field( ( @@ -389,7 +387,7 @@ def write_fields( pre=pre_out, offset_provider={"Koff": dims.KDim}, ) - if _x == int(args.itime): # End timer on last iteration + if _x == int(args.itime): # End timer on last iteration end_time = time.time() elapsed_time = end_time - start_time From 876e81ffc4bde067de5f40cf0200edb52f875919 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 16 Sep 2025 12:49:47 +0200 Subject: [PATCH 08/23] remove qti field --- .../subgrid_scale_physics/muphys/core/thermo.py | 13 ++++++++++--- .../muphys/integration_tests/run_full_muphys.py | 16 ++++++---------- .../stencil_tests/test_saturation_adjustment.py | 15 +++++++++++++-- 3 files changed, 29 insertions(+), 15 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py index d26df338d7..059c3a4ccb 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py @@ -346,7 +346,9 @@ def _saturation_adjustment( qve: fa.CellKField[ta.wpfloat], qce: fa.CellKField[ta.wpfloat], qre: fa.CellKField[ta.wpfloat], - qti: fa.CellKField[ta.wpfloat], + qse: fa.CellKField[ta.wpfloat], + qie: fa.CellKField[ta.wpfloat], + qge: fa.CellKField[ta.wpfloat], rho: fa.CellKField[ta.wpfloat], ) -> tuple[ fa.CellKField[ta.wpfloat], @@ -371,6 +373,7 @@ def _saturation_adjustment( - Revised specific vapor content - Mask specifying where qce+qve less than holding capacity """ + qti = qse + qie + qge qt = qve + qce + qre + qti cvc = t_d.cvd * (1.0 - qt) + t_d.clw * qre + g_ct.ci * qti cv = cvc + t_d.cvv * qve + t_d.clw * qce @@ -405,11 +408,15 @@ def saturation_adjustment( qve: fa.CellKField[ta.wpfloat], # Specific humidity qce: fa.CellKField[ta.wpfloat], # Specific cloud water content qre: fa.CellKField[ta.wpfloat], # Specific rain water - qti: fa.CellKField[ta.wpfloat], # Specific mass of all ice species (total-ice) + qse: fa.CellKField[ta.wpfloat], # Specific snow water + qie: fa.CellKField[ta.wpfloat], # Specific ice water content + qge: fa.CellKField[ta.wpfloat], # Specific graupel water content rho: fa.CellKField[ta.wpfloat], # Density containing dry air and water constituents te_out: fa.CellKField[ta.wpfloat], # Temperature qve_out: fa.CellKField[ta.wpfloat], # Specific humidity qce_out: fa.CellKField[ta.wpfloat], # Specific cloud water content mask_out: fa.CellKField[bool], # Specific cloud water content ): - _saturation_adjustment(te, qve, qce, qre, qti, rho, out=(te_out, qve_out, qce_out, mask_out)) + _saturation_adjustment( + te, qve, qce, qre, qse, qie, qge, rho, out=(te_out, qve_out, qce_out, mask_out) + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 95af1590f1..1016c7bba0 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -199,14 +199,6 @@ def run_program(args, backend, data): np.transpose(data.qr[0, :, :]), allocator=backend, ) - qti = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qg[0, :, :] + data.qs[0, :, :] + data.qi[0, :, :]), - allocator=backend, - ) # Total ice rho = gtx.as_field( ( dims.CellDim, @@ -247,7 +239,9 @@ def run_program(args, backend, data): qve=qve, qce=qce, qre=qre, - qti=qti, # Total ice + qse=qse, + qie=qie, + qge=qge, rho=rho, te_out=t_out, # Temperature qve_out=qv_out, # Specific humidity @@ -294,7 +288,9 @@ def run_program(args, backend, data): qve=qve, qce=qce, qre=qre, - qti=qti, # Total ice + qse=qse, + qie=qie, + qge=qge, rho=rho, te_out=t_out, # Temperature qve_out=qv_out, # Specific humidity diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py index 1e91248a57..2c4f513fbc 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py @@ -6,7 +6,9 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations + from typing import TYPE_CHECKING + import numpy as np import pytest @@ -16,6 +18,7 @@ from icon4py.model.common.utils import data_allocation as data_alloc from icon4py.model.testing.stencil_tests import StencilTest + if TYPE_CHECKING: from icon4py.model.common.grid import base as base_grid @@ -32,7 +35,9 @@ def reference( qve: np.ndarray, qce: np.ndarray, qre: np.ndarray, - qti: np.ndarray, + qse: np.ndarray, + qie: np.ndarray, + qge: np.ndarray, rho: np.ndarray, **kwargs, ) -> dict: @@ -58,7 +63,13 @@ def input_data(self, grid: base_grid.Grid) -> dict: qre=data_alloc.constant_field( grid, 2.5939378002267028e-004, dims.CellDim, dims.KDim, dtype=wpfloat ), - qti=data_alloc.constant_field( + qse=data_alloc.constant_field( + grid, 1.0746937601645517e-005, dims.CellDim, dims.KDim, dtype=wpfloat + ), + qie=data_alloc.constant_field( + grid, 1.0746937601645517e-005, dims.CellDim, dims.KDim, dtype=wpfloat + ), + qge=data_alloc.constant_field( grid, 1.0746937601645517e-005, dims.CellDim, dims.KDim, dtype=wpfloat ), rho=data_alloc.constant_field( From d40b2157bb5e217aa023f5bd771c081c1afdafac Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 16 Sep 2025 13:08:45 +0200 Subject: [PATCH 09/23] fix out fields --- .../integration_tests/run_full_muphys.py | 103 +++++++++--------- 1 file changed, 53 insertions(+), 50 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 1016c7bba0..dbe94cc908 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -124,7 +124,9 @@ def calc_dz(ksize, z): return dz -def run_program(args, backend, data): +def run_program( + args, backend, data, t_out, qv_out, qc_out, qi_out, qr_out, qs_out, qg_out, pflx_out +): ksize = data.dz.shape[0] dz = gtx.as_field( @@ -208,6 +210,55 @@ def run_program(args, backend, data): allocator=backend, ) + pr_out = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.zeros((data.ncells, data.nlev)), + allocator=backend, + ) + ps_out = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.zeros((data.ncells, data.nlev)), + allocator=backend, + ) + pi_out = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.zeros((data.ncells, data.nlev)), + allocator=backend, + ) + pg_out = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.zeros((data.ncells, data.nlev)), + allocator=backend, + ) + pre_out = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.zeros((data.ncells, data.nlev)), + allocator=backend, + ) + mask_out = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + data.mask_out, + allocator=backend, + ) + saturation_adjustment_program = setup_program( backend=backend, program=saturation_adjustment, @@ -429,56 +480,8 @@ def write_fields( data.pflx_out, allocator=backend, ) -pr_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) -ps_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) -pi_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) -pg_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) -pre_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) -mask_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.mask_out, - allocator=backend, -) -run_program(args, backend, data) +run_program(args, backend, data, t_out, qv_out, qc_out, qi_out, qr_out, qs_out, qg_out, pflx_out) write_fields( args.output_file, From 8aff981a744999f516b42171a68cc43712d7403c Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 16 Sep 2025 14:15:11 +0200 Subject: [PATCH 10/23] edit run_graupel_only --- .../integration_tests/run_full_muphys.py | 1 - .../integration_tests/run_graupel_only.py | 191 ++++++++++-------- 2 files changed, 106 insertions(+), 86 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 640a3dacb0..4a3a00b3f7 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -407,7 +407,6 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() -backend = model_backends.BACKENDS["gtfn_gpu"] args = get_args() backend = model_backends.BACKENDS[args.backend] diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index fdd2009f2e..c79734fb3e 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -14,6 +14,8 @@ import gt4py.next as gtx import numpy as np +from icon4py.model.common.model_options import setup_program + try: from netCDF4 import Dataset @@ -160,14 +162,95 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() -backend = model_backends.BACKENDS["dace_gpu"] args = get_args() +backend = model_backends.BACKENDS[args.backend] set_lib_path(args.ldir) sys.setrecursionlimit(10**4) data = Data(args) +dz = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.dz[:, :]), + allocator=backend, +) +te = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.t[0, :, :]), + allocator=backend, +) +p = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.p[0, :, :]), + allocator=backend, +) +qse = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qs[0, :, :]), + allocator=backend, +) +qie = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qi[0, :, :]), + allocator=backend, +) +qge = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qg[0, :, :]), + allocator=backend, +) +qve = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qv[0, :, :]), + allocator=backend, +) +qce = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qc[0, :, :]), + allocator=backend, +) +qre = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.qr[0, :, :]), + allocator=backend, +) +rho = gtx.as_field( + ( + dims.CellDim, + dims.KDim, + ), + np.transpose(data.rho[0, :, :]), + allocator=backend, +) + t_out = gtx.as_field( ( dims.CellDim, @@ -283,93 +366,32 @@ def write_fields( ksize = data.dz.shape[0] +graupel_run_program = setup_program( + backend=backend, + program=graupel_run, + constant_args={}, + horizontal_sizes={}, + vertical_sizes={ + "last_lev": ksize - 1, + }, + offset_provider={"Koff": dims.KDim}, +) + for _x in range(int(args.itime) + 1): if _x == 1: # Only start timing second iteration start_time = time.time() - graupel_run = graupel_run.with_backend(backend) - graupel_run( - last_lev=ksize - 1, - dz=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.dz[:, :]), - allocator=backend, - ), - te=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.t[0, :, :]), - allocator=backend, - ), - p=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.p[0, :, :]), - allocator=backend, - ), - rho=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.rho[0, :, :]), - allocator=backend, - ), - qve=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qv[0, :, :]), - allocator=backend, - ), - qce=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qc[0, :, :]), - allocator=backend, - ), - qre=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qr[0, :, :]), - allocator=backend, - ), - qse=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qs[0, :, :]), - allocator=backend, - ), - qie=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qi[0, :, :]), - allocator=backend, - ), - qge=gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qg[0, :, :]), - allocator=backend, - ), + graupel_run_program( + dz=dz, + te=te, + p=p, + rho=rho, + qve=qve, + qce=qce, + qre=qre, + qse=qse, + qie=qie, + qge=qge, dt=args.dt, qnc=args.qnc, t_out=t_out, @@ -385,7 +407,6 @@ def write_fields( pi=pi_out, pg=pg_out, pre=pre_out, - offset_provider={"Koff": dims.KDim}, ) if _x == int(args.itime): # End timer on last iteration end_time = time.time() From ab9981fab57b7470273b951b74dd32b702334271 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 16 Sep 2025 14:36:35 +0200 Subject: [PATCH 11/23] add backend parser --- .../tests/muphys/integration_tests/run_graupel_only.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index c79734fb3e..70c314a3b4 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -43,6 +43,13 @@ def get_args(): help="output filename", default="output.nc", ) + parser.add_argument( + "-b", + metavar="backend", + dest="backend", + help="gt4py backend", + default="gtfn_cpu", + ) parser.add_argument("input_file", help="input data file") parser.add_argument("itime", help="time-index", nargs="?", default=0) parser.add_argument("dt", help="timestep", nargs="?", default=30.0) From 5a574e6516161f5026832bebe306148792401c40 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 16 Sep 2025 14:41:41 +0200 Subject: [PATCH 12/23] cleanup --- .../tests/muphys/integration_tests/run_full_muphys.py | 1 + .../tests/muphys/integration_tests/run_graupel_only.py | 9 +-------- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 4a3a00b3f7..dbe94cc908 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -407,6 +407,7 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() + args = get_args() backend = model_backends.BACKENDS[args.backend] diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index 70c314a3b4..d812acbe49 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -169,6 +169,7 @@ def write_fields( pre_gsp_var[:, :] = pre_gsp ncfile.close() + args = get_args() backend = model_backends.BACKENDS[args.backend] @@ -362,14 +363,6 @@ def write_fields( np.zeros((data.ncells, data.nlev)), allocator=backend, ) -mask_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.mask_out, - allocator=backend, -) ksize = data.dz.shape[0] From 37525924ee3a3da20004fb45ee70a00333296b11 Mon Sep 17 00:00:00 2001 From: Will Sawyer Date: Thu, 18 Sep 2025 12:36:41 +0200 Subject: [PATCH 13/23] Corrected inputs for qse, qie, qge so that qti=1.0746937601645517e-005 --- .../muphys/stencil_tests/test_saturation_adjustment.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py index 2c4f513fbc..91a3a17125 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py @@ -64,13 +64,13 @@ def input_data(self, grid: base_grid.Grid) -> dict: grid, 2.5939378002267028e-004, dims.CellDim, dims.KDim, dtype=wpfloat ), qse=data_alloc.constant_field( - grid, 1.0746937601645517e-005, dims.CellDim, dims.KDim, dtype=wpfloat + grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat ), qie=data_alloc.constant_field( - grid, 1.0746937601645517e-005, dims.CellDim, dims.KDim, dtype=wpfloat + grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat ), qge=data_alloc.constant_field( - grid, 1.0746937601645517e-005, dims.CellDim, dims.KDim, dtype=wpfloat + grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat ), rho=data_alloc.constant_field( grid, 1.1371657035251757, dims.CellDim, dims.KDim, dtype=wpfloat From d4432fc2ccd48ed287c1cb577f5a16f7025b5a5b Mon Sep 17 00:00:00 2001 From: Will Sawyer Date: Tue, 30 Sep 2025 14:34:39 +0200 Subject: [PATCH 14/23] Removed mask_out from saturation_adjustment -- not needed --- .../subgrid_scale_physics/muphys/core/thermo.py | 12 +++++------- .../stencil_tests/test_saturation_adjustment.py | 4 +--- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py index 059c3a4ccb..0eafdd269b 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py @@ -354,7 +354,6 @@ def _saturation_adjustment( fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], - fa.CellKField[bool], ]: """ Compute the saturation adjustment which revises internal energy and water contents @@ -395,11 +394,11 @@ def _saturation_adjustment( # Is it possible to unify the where for all three outputs?? mask = qve + qce <= qx_hold - te = where((qve + qce <= qx_hold), Tx_hold, Tx) - qce = where((qve + qce <= qx_hold), 0.0, maximum(qve + qce - qx, 0.0)) - qve = where((qve + qce <= qx_hold), qve + qce, qx) + te = where(mask, Tx_hold, Tx) + qce = where(mask, 0.0, maximum(qve + qce - qx, 0.0)) + qve = where(mask, qve + qce, qx) - return te, qve, qce, mask + return te, qve, qce @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) @@ -415,8 +414,7 @@ def saturation_adjustment( te_out: fa.CellKField[ta.wpfloat], # Temperature qve_out: fa.CellKField[ta.wpfloat], # Specific humidity qce_out: fa.CellKField[ta.wpfloat], # Specific cloud water content - mask_out: fa.CellKField[bool], # Specific cloud water content ): _saturation_adjustment( - te, qve, qce, qre, qse, qie, qge, rho, out=(te_out, qve_out, qce_out, mask_out) + te, qve, qce, qre, qse, qie, qge, rho, out=(te_out, qve_out, qce_out) ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py index 91a3a17125..58b83e8f52 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py @@ -26,7 +26,7 @@ # @pytest.mark.embedded_only class TestSaturationAdjustment(StencilTest): PROGRAM = saturation_adjustment - OUTPUTS = ("te_out", "qve_out", "qce_out", "mask_out") + OUTPUTS = ("te_out", "qve_out", "qce_out") @staticmethod def reference( @@ -45,7 +45,6 @@ def reference( te_out=np.full(te.shape, 273.91226488486984), qve_out=np.full(te.shape, 4.4903852062454690e-003), qce_out=np.full(te.shape, 9.5724552280369163e-007), - mask_out=np.full(te.shape, False), ) @pytest.fixture @@ -78,5 +77,4 @@ def input_data(self, grid: base_grid.Grid) -> dict: te_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), qve_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), qce_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - mask_out=data_alloc.constant_field(grid, True, dims.CellDim, dims.KDim, dtype=bool), ) From d475f0bd8dcfb664122649362ed066bfb3acc6a0 Mon Sep 17 00:00:00 2001 From: Will Sawyer Date: Tue, 30 Sep 2025 14:48:25 +0200 Subject: [PATCH 15/23] Adjustments for performance benefit --- .../tests/muphys/integration_tests/run_full_muphys.py | 11 +++++------ .../muphys/integration_tests/run_graupel_only.py | 3 +++ 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index dbe94cc908..f6f9f65200 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -328,12 +328,6 @@ def run_program( pre=pre_out, ) - data.prr_gsp = np.transpose(pr_out[dims.KDim(ksize - 1)].asnumpy()) - data.prs_gsp = np.transpose(ps_out[dims.KDim(ksize - 1)].asnumpy()) - data.pri_gsp = np.transpose(pi_out[dims.KDim(ksize - 1)].asnumpy()) - data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) - data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) - saturation_adjustment_program( te=te, qve=qve, @@ -354,6 +348,11 @@ def run_program( elapsed_time = end_time - start_time print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") + data.prr_gsp = np.transpose(pr_out[dims.KDim(ksize - 1)].asnumpy()) + data.prs_gsp = np.transpose(ps_out[dims.KDim(ksize - 1)].asnumpy()) + data.pri_gsp = np.transpose(pi_out[dims.KDim(ksize - 1)].asnumpy()) + data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) + data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) def write_fields( diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index d812acbe49..1914741a3b 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -15,6 +15,7 @@ import numpy as np from icon4py.model.common.model_options import setup_program +from icon4py.model.common.utils import device_utils try: @@ -379,6 +380,7 @@ def write_fields( for _x in range(int(args.itime) + 1): if _x == 1: # Only start timing second iteration + device_utils.sync(backend) start_time = time.time() graupel_run_program( @@ -409,6 +411,7 @@ def write_fields( pre=pre_out, ) if _x == int(args.itime): # End timer on last iteration + device_utils.sync(backend) end_time = time.time() elapsed_time = end_time - start_time From e1dd9e481ceebe83526cb6405d9000586e8243bb Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 7 Nov 2025 18:11:12 +0100 Subject: [PATCH 16/23] add static compute domain --- .../muphys/core/thermo.py | 20 +++++++++++++++++-- .../muphys/implementations/graupel.py | 8 ++++++++ .../integration_tests/run_full_muphys.py | 17 +++++++++++++--- .../integration_tests/run_graupel_only.py | 7 ++++++- 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py index 0eafdd269b..31242b015c 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py @@ -9,7 +9,7 @@ from gt4py.next import exp, maximum, where from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.frozen import g_ct, t_d -from icon4py.model.common import field_type_aliases as fa, type_alias as ta +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta @gtx.field_operator @@ -414,7 +414,23 @@ def saturation_adjustment( te_out: fa.CellKField[ta.wpfloat], # Temperature qve_out: fa.CellKField[ta.wpfloat], # Specific humidity qce_out: fa.CellKField[ta.wpfloat], # Specific cloud water content + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, ): _saturation_adjustment( - te, qve, qce, qre, qse, qie, qge, rho, out=(te_out, qve_out, qce_out) + te, + qve, + qce, + qre, + qse, + qie, + qge, + rho, + out=(te_out, qve_out, qce_out), + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py index 2db883f645..31717fa2b8 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py @@ -505,6 +505,10 @@ def graupel_run( pi: fa.CellKField[ta.wpfloat], # Precipitation of ice pg: fa.CellKField[ta.wpfloat], # Precipitation of graupel pre: fa.CellKField[ta.wpfloat], # Precipitation of graupel + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, ): _graupel_run( last_lev, @@ -521,4 +525,8 @@ def graupel_run( dt, qnc, out=(t_out, qv_out, qc_out, qr_out, qs_out, qi_out, qg_out, pflx, pr, ps, pi, pg, pre), + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index f6f9f65200..4772b185d5 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -263,8 +263,14 @@ def run_program( backend=backend, program=saturation_adjustment, constant_args={}, - horizontal_sizes={}, - vertical_sizes={}, + horizontal_sizes={ + "horizontal_start": gtx.int32(0), + "horizontal_end": data.ncells, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(data.nlev), + }, offset_provider={"Koff": dims.KDim}, ) @@ -272,8 +278,13 @@ def run_program( backend=backend, program=graupel_run, constant_args={}, - horizontal_sizes={}, + horizontal_sizes={ + "horizontal_start": gtx.int32(0), + "horizontal_end": data.ncells, + }, vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(data.nlev), "last_lev": ksize - 1, }, offset_provider={"Koff": dims.KDim}, diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index 1914741a3b..dcf8a07f17 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -371,8 +371,13 @@ def write_fields( backend=backend, program=graupel_run, constant_args={}, - horizontal_sizes={}, + horizontal_sizes={ + "horizontal_start": gtx.int32(0), + "horizontal_end": data.ncells, + }, vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(data.nlev), "last_lev": ksize - 1, }, offset_provider={"Koff": dims.KDim}, From 9153482daa9c7bb20929781ce51d6d7723f65a49 Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 21 Nov 2025 13:45:16 +0100 Subject: [PATCH 17/23] use backend_like --- .../muphys/tests/muphys/integration_tests/run_full_muphys.py | 4 +++- .../muphys/tests/muphys/integration_tests/run_graupel_only.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py index 4772b185d5..1e2ae3f9b3 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py @@ -20,6 +20,7 @@ import netCDF4 import numpy as np +from icon4py.model.common import model_options from icon4py.model.common.model_options import setup_program @@ -419,7 +420,8 @@ def write_fields( args = get_args() -backend = model_backends.BACKENDS[args.backend] +backend_like = model_backends.BACKENDS[args.backend] +backend = model_options.customize_backend(None, backend_like) set_lib_path(args.ldir) sys.setrecursionlimit(10**4) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py index dcf8a07f17..59a1cb0a35 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py @@ -14,6 +14,7 @@ import gt4py.next as gtx import numpy as np +from icon4py.model.common import model_options from icon4py.model.common.model_options import setup_program from icon4py.model.common.utils import device_utils @@ -172,7 +173,8 @@ def write_fields( args = get_args() -backend = model_backends.BACKENDS[args.backend] +backend_like = model_backends.BACKENDS[args.backend] +backend = model_options.customize_backend(None, backend_like) set_lib_path(args.ldir) sys.setrecursionlimit(10**4) From 798ad60b9a2699689313a4c553115ced292739f2 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Thu, 27 Nov 2025 11:40:15 +0100 Subject: [PATCH 18/23] muphys: Refactor graupel_only driver and add integration test (#958) Clean up the graupel_only driver, and create an integration test to run through pytest. Co-authored-by: Will Sawyer --- .../muphys/driver/__init__.py | 7 + .../muphys/driver}/run_full_muphys.py | 10 +- .../muphys/driver/run_graupel_only.py | 296 ++++++++++++ .../muphys/driver/utils.py | 23 + .../integration_tests/run_graupel_only.py | 450 ------------------ .../integration_tests/test_graupel_only.py | 143 ++++++ .../icon4py/model/testing/data_handling.py | 23 +- .../model/testing/fixtures/datatest.py | 25 +- 8 files changed, 497 insertions(+), 480 deletions(-) create mode 100644 model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/__init__.py rename model/atmosphere/subgrid_scale_physics/muphys/{tests/muphys/integration_tests => src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver}/run_full_muphys.py (98%) create mode 100755 model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py create mode 100644 model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/utils.py delete mode 100755 model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py create mode 100644 model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/__init__.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/__init__.py new file mode 100644 index 0000000000..de9850de36 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/__init__.py @@ -0,0 +1,7 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py similarity index 98% rename from model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py rename to model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py index 1e2ae3f9b3..c9d57b3b60 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py @@ -12,6 +12,8 @@ WORK IN PROGRESS!!!! Do not try to run this. """ +# TODO refactor similar to run_graupel_only + import argparse import sys import time @@ -385,10 +387,10 @@ def write_fields( pflx, pre_gsp, ): - ncfile = Dataset(output_filename, mode="w") - ncells = ncfile.createDimension("ncells", ncell) - height = ncfile.createDimension("height", nlev) - height1 = ncfile.createDimension("height1", nlev+1) + ncfile = Dataset(output_filename, mode="w") + ncells = ncfile.createDimension("ncells", ncell) + height = ncfile.createDimension("height", nlev) + height1 = ncfile.createDimension("height1", nlev + 1) ta_var = ncfile.createVariable("ta", np.double, ("height", "ncells")) hus_var = ncfile.createVariable("hus", np.double, ("height", "ncells")) clw_var = ncfile.createVariable("clw", np.double, ("height", "ncells")) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py new file mode 100755 index 0000000000..1c67c86c21 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py @@ -0,0 +1,296 @@ +#!/usr/bin/env python +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import argparse +import dataclasses +import functools +import pathlib +import time + +import netCDF4 +import numpy as np +from gt4py import next as gtx +from gt4py.next import typing as gtx_typing + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import utils +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations import graupel +from icon4py.model.common import dimension as dims, model_backends, model_options +from icon4py.model.common.utils import device_utils + + +# TODO(havogt): make similar to icon4py driver structure + + +def get_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "-o", metavar="output_file", dest="output_file", help="output filename", default="output.nc" + ) + parser.add_argument( + "-b", metavar="backend", dest="backend", help="gt4py backend", default="gtfn_cpu" + ) + parser.add_argument("input_file", help="input data file") + parser.add_argument("itime", help="time-index", nargs="?", default=0) + parser.add_argument("dt", help="timestep", nargs="?", default=30.0) + parser.add_argument("qnc", help="Water number concentration", nargs="?", default=100.0) + + return parser.parse_args() + + +# TODO double check the sizes of pxxx vars (should be nlev+1?) + + +def _calc_dz(z: np.ndarray) -> np.ndarray: + ksize = z.shape[0] + dz = np.zeros(z.shape, np.float64) + zh = 1.5 * z[ksize - 1, :] - 0.5 * z[ksize - 2, :] + for k in range(ksize - 1, -1, -1): + zh_new = 2.0 * z[k, :] - zh + dz[k, :] = -zh + zh_new + zh = zh_new + return dz + + +def _as_field_from_nc( + dataset: netCDF4.Dataset, + allocator: gtx_typing.FieldBufferAllocationUtil, + varname: str, + optional: bool = False, + dtype: np.dtype | None = None, +) -> gtx.Field[dims.CellDim, dims.KDim] | None: + if optional and varname not in dataset.variables: + return None + + var = dataset.variables[varname] + if var.dimensions[0] == "time": + var = var[0, :, :] + data = np.transpose(var) + if dtype is not None: + data = data.astype(dtype) + return gtx.as_field( + (dims.CellDim, dims.KDim), + data, + allocator=allocator, + ) + + +def _field_to_nc( + dataset: netCDF4.Dataset, + dims: tuple[str, str], + varname: str, + field: gtx.Field[dims.CellDim, dims.KDim], + dtype: np.dtype = np.float64, +) -> None: + var = dataset.createVariable(varname, dtype, dims) + var[...] = field.asnumpy().transpose() + + +@dataclasses.dataclass +class GraupelInput: + ncells: int + nlev: int + dz: gtx.Field[dims.CellDim, dims.KDim] + p: gtx.Field[dims.CellDim, dims.KDim] + rho: gtx.Field[dims.CellDim, dims.KDim] + t: gtx.Field[dims.CellDim, dims.KDim] + qv: gtx.Field[dims.CellDim, dims.KDim] + qc: gtx.Field[dims.CellDim, dims.KDim] + qi: gtx.Field[dims.CellDim, dims.KDim] + qr: gtx.Field[dims.CellDim, dims.KDim] + qs: gtx.Field[dims.CellDim, dims.KDim] + qg: gtx.Field[dims.CellDim, dims.KDim] + + @classmethod + def load( + cls, filename: pathlib.Path | str, allocator: gtx_typing.FieldBufferAllocationUtil + ) -> None: + with netCDF4.Dataset(filename, mode="r") as ncfile: + try: + ncells = len(ncfile.dimensions["cell"]) + except KeyError: + ncells = len(ncfile.dimensions["ncells"]) + + nlev = len(ncfile.dimensions["height"]) + + dz = _calc_dz(ncfile.variables["zg"]) + + field_from_nc = functools.partial( + _as_field_from_nc, ncfile, allocator, dtype=np.float64 + ) + return cls( + ncells=ncells, + nlev=nlev, + dz=gtx.as_field((dims.CellDim, dims.KDim), np.transpose(dz), allocator=allocator), + t=field_from_nc("ta"), + p=field_from_nc("pfull"), + qs=field_from_nc("qs"), + qi=field_from_nc("cli"), + qg=field_from_nc("qg"), + qv=field_from_nc("hus"), + qc=field_from_nc("clw"), + qr=field_from_nc("qr"), + rho=field_from_nc("rho"), + ) + + +@dataclasses.dataclass +class GraupelOutput: + t: gtx.Field[dims.CellDim, dims.KDim] + qv: gtx.Field[dims.CellDim, dims.KDim] + qc: gtx.Field[dims.CellDim, dims.KDim] + qi: gtx.Field[dims.CellDim, dims.KDim] + qr: gtx.Field[dims.CellDim, dims.KDim] + qs: gtx.Field[dims.CellDim, dims.KDim] + qg: gtx.Field[dims.CellDim, dims.KDim] + + pflx: gtx.Field[dims.CellDim, dims.KDim] | None + pr: gtx.Field[dims.CellDim, dims.KDim] | None + ps: gtx.Field[dims.CellDim, dims.KDim] | None + pi: gtx.Field[dims.CellDim, dims.KDim] | None + pg: gtx.Field[dims.CellDim, dims.KDim] | None + pre: gtx.Field[dims.CellDim, dims.KDim] | None + + @classmethod + def allocate(cls, allocator: gtx_typing.FieldBufferAllocationUtil, domain: gtx.Domain): + zeros = functools.partial(gtx.zeros, domain=domain, allocator=allocator) + # TODO +1 size fields? + return cls(**{field.name: zeros() for field in dataclasses.fields(cls)}) + + @classmethod + def load(cls, filename: pathlib.Path | str, allocator: gtx_typing.FieldBufferAllocationUtil): + with netCDF4.Dataset(filename, mode="r") as ncfile: + field_from_nc = functools.partial(_as_field_from_nc, ncfile, allocator) + return cls( + t=field_from_nc("ta"), + qv=field_from_nc("hus"), + qc=field_from_nc("clw"), + qi=field_from_nc("cli"), + qr=field_from_nc("qr"), + qs=field_from_nc("qs"), + qg=field_from_nc("qg"), + pflx=field_from_nc("pflx", optional=True), + pr=field_from_nc("prr_gsp", optional=True), + ps=field_from_nc("prs_gsp", optional=True), + pi=field_from_nc("pri_gsp", optional=True), + pg=field_from_nc("prg_gsp", optional=True), + pre=field_from_nc("pre_gsp", optional=True), + ) + + def write(self, filename: pathlib.Path | str): + ncells = self.t.shape[0] + nlev = self.t.shape[1] + + with netCDF4.Dataset(filename, mode="w") as ncfile: + ncfile.createDimension("ncells", ncells) + ncfile.createDimension("height", nlev) + + write_height_field = functools.partial( + _field_to_nc, ncfile, ("height", "ncells"), dtype=np.float64 + ) + + write_height_field("ta", self.t) + write_height_field("hus", self.qv) + write_height_field("clw", self.qc) + write_height_field("cli", self.qi) + write_height_field("qr", self.qr) + write_height_field("qs", self.qs) + write_height_field("qg", self.qg) + if self.pflx is not None: + write_height_field("pflx", self.pflx) + if self.pr is not None: + write_height_field("prr_gsp", self.pr) + if self.ps is not None: + write_height_field("prs_gsp", self.ps) # TODO + if self.pi is not None: + write_height_field("pri_gsp", self.pi) # TODO + if self.pg is not None: + write_height_field("prg_gsp", self.pg) # TODO + if self.pre is not None: + write_height_field("pre_gsp", self.pre) # TODO + + +def setup_graupel(inp: GraupelInput, dt: float, qnc: float, backend: model_backends.BackendLike): + with utils.recursion_limit(10**4): # TODO thread safe? + graupel_run_program = model_options.setup_program( + backend=backend, + program=graupel.graupel_run, + constant_args={"dt": dt, "qnc": qnc}, + horizontal_sizes={ + "horizontal_start": gtx.int32(0), + "horizontal_end": inp.ncells, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(inp.nlev), + "last_lev": gtx.int32(inp.nlev - 1), + }, + offset_provider={"Koff": dims.KDim}, + ) + gtx.wait_for_compilation() + return graupel_run_program + + +def main(): + args = get_args() + + backend = model_backends.BACKENDS[args.backend] + allocator = model_backends.get_allocator(backend) + + inp = GraupelInput.load(filename=pathlib.Path(args.input_file), allocator=allocator) + out = GraupelOutput.allocate( + domain=gtx.domain({dims.CellDim: inp.ncells, dims.KDim: inp.nlev}), allocator=allocator + ) + + graupel_run_program = setup_graupel(inp, dt=args.dt, qnc=args.qnc, backend=backend) + + start_time = None + for _x in range(int(args.itime) + 1): + if _x == 1: # Only start timing second iteration + device_utils.sync(backend) + start_time = time.time() + + graupel_run_program( + dz=inp.dz, + te=inp.t, + p=inp.p, + rho=inp.rho, + qve=inp.qv, + qce=inp.qc, + qre=inp.qr, + qse=inp.qs, + qie=inp.qi, + qge=inp.qg, + t_out=out.t, + qv_out=out.qv, + qc_out=out.qc, + qr_out=out.qr, + qs_out=out.qs, + qi_out=out.qi, + qg_out=out.qg, + pflx=out.pflx, + pr=out.pr, + ps=out.ps, + pi=out.pi, + pg=out.pg, + pre=out.pre, + ) + device_utils.sync(backend) + end_time = time.time() + + if start_time is not None: + elapsed_time = end_time - start_time + print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") + + out.write(args.output_file) + + +if __name__ == "__main__": + main() diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/utils.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/utils.py new file mode 100644 index 0000000000..ffe43ba31d --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/utils.py @@ -0,0 +1,23 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import contextlib +import sys +from collections.abc import Generator + + +@contextlib.contextmanager +def recursion_limit(limit: int) -> Generator[None, None, None]: + original_limit = sys.getrecursionlimit() + sys.setrecursionlimit(limit) + try: + yield + finally: + sys.setrecursionlimit(original_limit) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py deleted file mode 100755 index 59a1cb0a35..0000000000 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/run_graupel_only.py +++ /dev/null @@ -1,450 +0,0 @@ -#!/usr/bin/env python -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -import argparse -import sys -import time - -import gt4py.next as gtx -import numpy as np - -from icon4py.model.common import model_options -from icon4py.model.common.model_options import setup_program -from icon4py.model.common.utils import device_utils - - -try: - from netCDF4 import Dataset -except ImportError: - print("Netcdf not installed") - sys.exit() - - -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations.graupel import ( - graupel_run, -) -from icon4py.model.common import dimension as dims, model_backends - - -def set_lib_path(lib_dir): - sys.path.append(lib_dir) - - -def get_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "-o", - metavar="output_file", - dest="output_file", - help="output filename", - default="output.nc", - ) - parser.add_argument( - "-b", - metavar="backend", - dest="backend", - help="gt4py backend", - default="gtfn_cpu", - ) - parser.add_argument("input_file", help="input data file") - parser.add_argument("itime", help="time-index", nargs="?", default=0) - parser.add_argument("dt", help="timestep", nargs="?", default=30.0) - parser.add_argument("qnc", help="Water number concentration", nargs="?", default=100.0) - parser.add_argument( - "-ldir", - metavar="lib_dir", - dest="ldir", - help="directory with py_graupel shared lib", - default="build/lib64", - ) - parser.add_argument( - "-with_sat_adj", - action="store_true", - ) - - return parser.parse_args() - - -class Data: - def __init__(self, args): - nc = Dataset(args.input_file) - # intent(in) variables: - try: - self.ncells = len(nc.dimensions["cell"]) - except KeyError: - self.ncells = len(nc.dimensions["ncells"]) - - self.nlev = len(nc.dimensions["height"]) - self.z = nc.variables["zg"][:, :].astype(np.float64) - self.p = nc.variables["pfull"][:, :].astype(np.float64) - self.rho = nc.variables["rho"][:, :].astype(np.float64) - # intent(inout) variables: - self.t = nc.variables["ta"][:, :].astype(np.float64) # inout - self.qv = nc.variables["hus"][:, :].astype(np.float64) # inout - self.qc = nc.variables["clw"][:, :].astype(np.float64) # inout - self.qi = nc.variables["cli"][:, :].astype(np.float64) # inout - self.qr = nc.variables["qr"][:, :].astype(np.float64) # inout - self.qs = nc.variables["qs"][:, :].astype(np.float64) # inout - self.qg = nc.variables["qg"][:, :].astype(np.float64) # inout - # intent(out) variables: - self.t_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qv_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qc_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qi_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qr_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qs_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qg_out = np.zeros((self.ncells, self.nlev), np.float64) - self.pflx_out = np.zeros((self.ncells, self.nlev), np.float64) - self.prr_gsp = np.zeros(self.ncells, np.float64) - self.pri_gsp = np.zeros(self.ncells, np.float64) - self.prs_gsp = np.zeros(self.ncells, np.float64) - self.prg_gsp = np.zeros(self.ncells, np.float64) - self.pre_gsp = np.zeros(self.ncells, np.float64) - self.dz = calc_dz(self.nlev, self.z) - self.mask_out = np.full((self.ncells, self.nlev), True) - - -def calc_dz(ksize, z): - dz = np.zeros(z.shape, np.float64) - zh = 1.5 * z[ksize - 1, :] - 0.5 * z[ksize - 2, :] - for k in range(ksize - 1, -1, -1): - zh_new = 2.0 * z[k, :] - zh - dz[k, :] = -zh + zh_new - zh = zh_new - return dz - - -def write_fields( - output_filename, - ncell, - nlev, - t, - qv, - qc, - qi, - qr, - qs, - qg, - prr_gsp, - prs_gsp, - pri_gsp, - prg_gsp, - pflx, - pre_gsp, -): - ncfile = Dataset(output_filename, mode="w") - ncells = ncfile.createDimension("ncells", ncell) - height = ncfile.createDimension("height", nlev) - height1 = ncfile.createDimension("height1", nlev+1) - ta_var = ncfile.createVariable("ta", np.double, ("height", "ncells")) - hus_var = ncfile.createVariable("hus", np.double, ("height", "ncells")) - clw_var = ncfile.createVariable("clw", np.double, ("height", "ncells")) - cli_var = ncfile.createVariable("cli", np.double, ("height", "ncells")) - qr_var = ncfile.createVariable("qr", np.double, ("height", "ncells")) - qs_var = ncfile.createVariable("qs", np.double, ("height", "ncells")) - qg_var = ncfile.createVariable("qg", np.double, ("height", "ncells")) - pflx_var = ncfile.createVariable("pflx", np.double, ("height", "ncells")) - prr_gsp_var = ncfile.createVariable("prr_gsp", np.double, ("height1", "ncells")) - prs_gsp_var = ncfile.createVariable("prs_gsp", np.double, ("height1", "ncells")) - pri_gsp_var = ncfile.createVariable("pri_gsp", np.double, ("height1", "ncells")) - prg_gsp_var = ncfile.createVariable("prg_gsp", np.double, ("height1", "ncells")) - pre_gsp_var = ncfile.createVariable("pre_gsp", np.double, ("height1", "ncells")) - - ta_var[:, :] = t - hus_var[:, :] = qv - clw_var[:, :] = qc - cli_var[:, :] = qi - qr_var[:, :] = qr - qs_var[:, :] = qs - qg_var[:, :] = qg - pflx_var[:, :] = pflx - prr_gsp_var[:, :] = prr_gsp - prs_gsp_var[:, :] = prs_gsp - pri_gsp_var[:, :] = pri_gsp - prg_gsp_var[:, :] = prg_gsp - pre_gsp_var[:, :] = pre_gsp - ncfile.close() - - -args = get_args() -backend_like = model_backends.BACKENDS[args.backend] -backend = model_options.customize_backend(None, backend_like) - -set_lib_path(args.ldir) -sys.setrecursionlimit(10**4) - -data = Data(args) - -dz = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.dz[:, :]), - allocator=backend, -) -te = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.t[0, :, :]), - allocator=backend, -) -p = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.p[0, :, :]), - allocator=backend, -) -qse = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qs[0, :, :]), - allocator=backend, -) -qie = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qi[0, :, :]), - allocator=backend, -) -qge = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qg[0, :, :]), - allocator=backend, -) -qve = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qv[0, :, :]), - allocator=backend, -) -qce = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qc[0, :, :]), - allocator=backend, -) -qre = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qr[0, :, :]), - allocator=backend, -) -rho = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.rho[0, :, :]), - allocator=backend, -) - -t_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.t_out, - allocator=backend, -) -qv_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qv_out, - allocator=backend, -) -qc_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qc_out, - allocator=backend, -) -qr_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qr_out, - allocator=backend, -) -qs_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qs_out, - allocator=backend, -) -qi_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qi_out, - allocator=backend, -) -qg_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qg_out, - allocator=backend, -) -pflx_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.pflx_out, - allocator=backend, -) -pr_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) -ps_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) -pi_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) -pg_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) -pre_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, -) - -ksize = data.dz.shape[0] - -graupel_run_program = setup_program( - backend=backend, - program=graupel_run, - constant_args={}, - horizontal_sizes={ - "horizontal_start": gtx.int32(0), - "horizontal_end": data.ncells, - }, - vertical_sizes={ - "vertical_start": gtx.int32(0), - "vertical_end": gtx.int32(data.nlev), - "last_lev": ksize - 1, - }, - offset_provider={"Koff": dims.KDim}, -) - -for _x in range(int(args.itime) + 1): - if _x == 1: # Only start timing second iteration - device_utils.sync(backend) - start_time = time.time() - - graupel_run_program( - dz=dz, - te=te, - p=p, - rho=rho, - qve=qve, - qce=qce, - qre=qre, - qse=qse, - qie=qie, - qge=qge, - dt=args.dt, - qnc=args.qnc, - t_out=t_out, - qv_out=qv_out, - qc_out=qc_out, - qr_out=qr_out, - qs_out=qs_out, - qi_out=qi_out, - qg_out=qg_out, - pflx=pflx_out, - pr=pr_out, - ps=ps_out, - pi=pi_out, - pg=pg_out, - pre=pre_out, - ) - if _x == int(args.itime): # End timer on last iteration - device_utils.sync(backend) - end_time = time.time() - -elapsed_time = end_time - start_time -print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") - -data.prr_gsp = np.transpose(pr_out[dims.KDim(ksize - 1)].asnumpy()) -data.prs_gsp = np.transpose(ps_out[dims.KDim(ksize - 1)].asnumpy()) -data.pri_gsp = np.transpose(pi_out[dims.KDim(ksize - 1)].asnumpy()) -data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) -data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) - -write_fields( - args.output_file, - data.ncells, - data.nlev, - t=np.transpose(t_out.asnumpy()), - qv=np.transpose(qv_out.asnumpy()), - qc=np.transpose(qc_out.asnumpy()), - qi=np.transpose(qi_out.asnumpy()), - qr=np.transpose(qr_out.asnumpy()), - qs=np.transpose(qs_out.asnumpy()), - qg=np.transpose(qg_out.asnumpy()), - prr_gsp=data.prr_gsp, - pri_gsp=data.pri_gsp, - prs_gsp=data.prs_gsp, - prg_gsp=data.prg_gsp, - pflx=np.transpose(pflx_out.asnumpy()), - pre_gsp=data.pre_gsp, -) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py new file mode 100644 index 0000000000..11c1be1ff2 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py @@ -0,0 +1,143 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import dataclasses +import pathlib +from typing import Final + +import numpy as np +import pytest +from gt4py import next as gtx + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import run_graupel_only +from icon4py.model.common import dimension as dims, model_backends +from icon4py.model.testing import data_handling, definitions as testing_defs +from icon4py.model.testing.fixtures.datatest import backend_like + + +def _path_to_experiment_testdata(experiment: MuphysGraupelExperiment) -> pathlib.Path: + return testing_defs.get_test_data_root_path() / "muphys_graupel_data" / experiment.name + + +@dataclasses.dataclass(frozen=True) +class MuphysGraupelExperiment: + name: str + uri: str + dtype: np.dtype + dt: float = 30.0 + qnc: float = 100.0 + + @property + def input_file(self) -> pathlib.Path: + return _path_to_experiment_testdata(self) / "input.nc" + + @property + def reference_file(self) -> pathlib.Path: + return _path_to_experiment_testdata(self) / "reference.nc" + + def __str__(self): + return self.name + + +class Experiments: + # TODO currently on havogt's polybox + MINI: Final = MuphysGraupelExperiment( + name="mini", + uri="https://polybox.ethz.ch/index.php/s/55oHBDxS2SiqAGN/download/mini.tar.gz", + dtype=np.float32, + ) + TINY: Final = MuphysGraupelExperiment( + name="tiny", + uri="https://polybox.ethz.ch/index.php/s/5Ceop3iaWkbc7gf/download/tiny.tar.gz", + dtype=np.float64, + ) + R2B05: Final = MuphysGraupelExperiment( + name="R2B05", + uri="https://polybox.ethz.ch/index.php/s/RBib8rFSEd7Eomo/download/R2B05.tar.gz", + dtype=np.float32, + ) + + +@pytest.fixture(autouse=True) +def download_test_data(experiment: MuphysGraupelExperiment) -> None: + """Downloads test data for an experiment (implicit fixture).""" + data_handling.download_test_data(_path_to_experiment_testdata(experiment), uri=experiment.uri) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "experiment", + [ + Experiments.MINI, + Experiments.TINY, + Experiments.R2B05, + ], + ids=lambda exp: exp.name, +) +def test_graupel_only( + backend_like: model_backends.BackendLike, + experiment: MuphysGraupelExperiment, +) -> None: + inp = run_graupel_only.GraupelInput.load( + filename=experiment.input_file, allocator=model_backends.get_allocator(backend_like) + ) + + graupel_run_program = run_graupel_only.setup_graupel( + inp, dt=experiment.dt, qnc=experiment.qnc, backend=backend_like + ) + + out = run_graupel_only.GraupelOutput.allocate( + allocator=model_backends.get_allocator(backend_like), + domain=gtx.domain({dims.CellDim: inp.ncells, dims.KDim: inp.nlev}), + ) + + graupel_run_program( + dz=inp.dz, + te=inp.t, + p=inp.p, + rho=inp.rho, + qve=inp.qv, + qce=inp.qc, + qre=inp.qr, + qse=inp.qs, + qie=inp.qi, + qge=inp.qg, + t_out=out.t, + qv_out=out.qv, + qc_out=out.qc, + qr_out=out.qr, + qs_out=out.qs, + qi_out=out.qi, + qg_out=out.qg, + pflx=out.pflx, + pr=out.pr, + ps=out.ps, + pi=out.pi, + pg=out.pg, + pre=out.pre, + ) + + ref = run_graupel_only.GraupelOutput.load( + filename=experiment.reference_file, + allocator=model_backends.get_allocator(backend_like), + ) + + # TODO check tolerances + rtol = 1e-14 if experiment.dtype == np.float64 else 1e-7 + atol = 1e-16 if experiment.dtype == np.float64 else 1e-8 + # TODO we run the float32 input experiments with float64 + + np.testing.assert_allclose(ref.qv.asnumpy(), out.qv.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qc.asnumpy(), out.qc.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qi.asnumpy(), out.qi.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qr.asnumpy(), out.qr.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qs.asnumpy(), out.qs.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qg.asnumpy(), out.qg.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.t.asnumpy(), out.t.asnumpy(), atol=atol, rtol=rtol) diff --git a/model/testing/src/icon4py/model/testing/data_handling.py b/model/testing/src/icon4py/model/testing/data_handling.py index 9ecf932335..9624c64839 100644 --- a/model/testing/src/icon4py/model/testing/data_handling.py +++ b/model/testing/src/icon4py/model/testing/data_handling.py @@ -6,11 +6,13 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +import pathlib import tarfile -from pathlib import Path +from icon4py.model.testing import config, locking -def download_and_extract(uri: str, dst: Path, data_file: str = "downloaded.tar.gz") -> None: + +def download_and_extract(uri: str, dst: pathlib.Path, data_file: str = "downloaded.tar.gz") -> None: """ Download data archive from remote server. @@ -31,4 +33,19 @@ def download_and_extract(uri: str, dst: Path, data_file: str = "downloaded.tar.g raise OSError(f"{data_file} needs to be a valid tar file") with tarfile.open(data_file, mode="r:*") as tf: tf.extractall(path=dst) - Path(data_file).unlink(missing_ok=True) + pathlib.Path(data_file).unlink(missing_ok=True) + + +def download_test_data(dst: pathlib.Path, uri: str) -> None: + if config.ENABLE_TESTDATA_DOWNLOAD: + # We create and lock the *parent* directory as we later check for existence of `dst`. + dst.parent.mkdir(parents=True, exist_ok=True) + with locking.lock(dst.parent): + if not dst.exists(): + download_and_extract(uri, dst) + else: + # If test data download is disabled, we check if the directory exists + # without locking. We assume the location is managed by the user + # and avoid locking shared directories (e.g. on CI). + if not dst.exists(): + raise RuntimeError(f"Test data {dst} does not exist, and downloading is disabled.") diff --git a/model/testing/src/icon4py/model/testing/fixtures/datatest.py b/model/testing/src/icon4py/model/testing/fixtures/datatest.py index 123884f464..09b9f0836b 100644 --- a/model/testing/src/icon4py/model/testing/fixtures/datatest.py +++ b/model/testing/src/icon4py/model/testing/fixtures/datatest.py @@ -17,13 +17,7 @@ from icon4py.model.common import model_backends, model_options from icon4py.model.common.constants import RayleighType from icon4py.model.common.grid import base as base_grid -from icon4py.model.testing import ( - config, - data_handling as data, - datatest_utils as dt_utils, - definitions, - locking, -) +from icon4py.model.testing import data_handling as data, datatest_utils as dt_utils, definitions if TYPE_CHECKING: @@ -121,22 +115,7 @@ def _download_ser_data( try: destination_path = dt_utils.get_datapath_for_experiment(_ranked_data_path, _experiment) uri = _experiment.partitioned_data[comm_size] - - data_file = _ranked_data_path.joinpath(f"{_experiment.name}_mpitask{comm_size}.tar.gz").name - _ranked_data_path.mkdir(parents=True, exist_ok=True) - if config.ENABLE_TESTDATA_DOWNLOAD: - with locking.lock(_ranked_data_path): - # Note: if the lock would be created for `destination_path` it would always exist... - if not destination_path.exists(): - data.download_and_extract(uri, _ranked_data_path, data_file) - else: - # If test data download is disabled, we check if the directory exists - # without locking. We assume the location is managed by the user - # and avoid locking shared directories (e.g. on CI). - if not destination_path.exists(): - raise RuntimeError( - f"Serialization data {data_file} does not exist, and downloading is disabled." - ) + data.download_test_data(destination_path, uri) except KeyError as err: raise RuntimeError( f"No data for communicator of size {comm_size} exists, use 1, 2 or 4" From b9349217725602b8ae73635da300da9c5ab21dbc Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 5 Dec 2025 15:50:05 +0100 Subject: [PATCH 19/23] [draft] Muphys bug fix graupel refactoring hannes (#961) Co-authored-by: Will Sawyer Co-authored-by: Will Sawyer --- .../muphys/pyproject.toml | 1 + .../muphys/core/definitions.py | 20 + .../muphys/core/saturation_adjustment.py | 92 +++ .../muphys/core/saturation_adjustment2.py | 134 ---- .../muphys/core/thermo.py | 100 +-- .../muphys/driver/common.py | 216 ++++++ .../muphys/driver/run_full_muphys.py | 650 +++++------------- .../muphys/driver/run_graupel_only.py | 204 +----- .../muphys/implementations/graupel.py | 353 +++++----- .../muphys/implementations/muphys.py | 101 +++ .../integration_tests/test_full_muphys.py | 139 ++++ .../integration_tests/test_graupel_only.py | 29 +- .../test_saturation_adjustment.py | 53 +- .../test_saturation_adjustment2.py | 81 --- .../icon4py/model/testing/stencil_tests.py | 5 +- 15 files changed, 953 insertions(+), 1225 deletions(-) create mode 100644 model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/definitions.py create mode 100644 model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/saturation_adjustment.py delete mode 100644 model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/saturation_adjustment2.py create mode 100644 model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/common.py create mode 100644 model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/muphys.py create mode 100644 model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_full_muphys.py delete mode 100644 model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment2.py diff --git a/model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml b/model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml index 09f978bb96..65666169d3 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml +++ b/model/atmosphere/subgrid_scale_physics/muphys/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ # workspace members "icon4py-common[io]>=0.0.6", # external dependencies + "numpy>=1.23.3", "gt4py==1.1.1", "packaging>=20.0" ] diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/definitions.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/definitions.py new file mode 100644 index 0000000000..c1762c55f0 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/definitions.py @@ -0,0 +1,20 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from typing import NamedTuple + +from icon4py.model.common import field_type_aliases as fa, type_alias as ta + + +class Q(NamedTuple): + v: fa.CellKField[ta.wpfloat] # Specific humidity + c: fa.CellKField[ta.wpfloat] # Specific cloud water content + r: fa.CellKField[ta.wpfloat] # Specific rain water + s: fa.CellKField[ta.wpfloat] # Specific snow water + i: fa.CellKField[ta.wpfloat] # Specific ice water content + g: fa.CellKField[ta.wpfloat] # Specific graupel water content diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/saturation_adjustment.py new file mode 100644 index 0000000000..a308138790 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/saturation_adjustment.py @@ -0,0 +1,92 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause +import gt4py.next as gtx +from gt4py.next import maximum, where + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.frozen import g_ct, t_d +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import Q +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import ( + _newton_raphson, + _qsat_rho, +) +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta + + +@gtx.field_operator +def _saturation_adjustment( + te: fa.CellKField[ta.wpfloat], rho: fa.CellKField[ta.wpfloat], q_in: Q +) -> tuple[ + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], +]: + """ + Compute the saturation adjustment which revises internal energy and water contents + + Args: + Tx: Temperature + rho: Density containing dry air and water constituents + Q: Class with humidity, cloud, rain, snow, ice and graupel water + + Result: Tuple containing + - Revised temperature + - Revised specific cloud water content + - Revised specific vapor content + """ + qti = q_in.s + q_in.i + q_in.g + qt = q_in.v + q_in.c + q_in.r + qti + cvc = t_d.cvd * (1.0 - qt) + t_d.clw * q_in.r + g_ct.ci * qti + cv = cvc + t_d.cvv * q_in.v + t_d.clw * q_in.c + ue = cv * te - q_in.c * g_ct.lvc + Tx_hold = ue / (cv + q_in.c * (t_d.cvv - t_d.clw)) + qx_hold = _qsat_rho(Tx_hold, rho) + + Tx = te + # Newton-Raphson iteration: 6 times the same operations + Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue) + Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue) + Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue) + Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue) + Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue) + Tx = _newton_raphson(Tx, rho, q_in.v, q_in.c, cvc, ue) + + # At this point we hope Tx has converged + qx = _qsat_rho(Tx, rho) + + # Is it possible to unify the where for all three outputs?? + mask = q_in.v + q_in.c <= qx_hold + te = where(mask, Tx_hold, Tx) + qce = where(mask, 0.0, maximum(q_in.v + q_in.c - qx, 0.0)) + qve = where(mask, q_in.v + q_in.c, qx) + + return te, qve, qce + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def saturation_adjustment( + te: fa.CellKField[ta.wpfloat], # Temperature + rho: fa.CellKField[ta.wpfloat], # Density containing dry air and water constituents + q_in: Q, # Class with humidity, cloud, rain, snow, ice and graupel water + te_out: fa.CellKField[ta.wpfloat], # Temperature + qve_out: fa.CellKField[ta.wpfloat], # Specific humidity + qce_out: fa.CellKField[ta.wpfloat], # Specific cloud water content + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _saturation_adjustment( + te, + rho, + q_in, + out=(te_out, qve_out, qce_out), + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/saturation_adjustment2.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/saturation_adjustment2.py deleted file mode 100644 index 353ff20603..0000000000 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/saturation_adjustment2.py +++ /dev/null @@ -1,134 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause - -""" -This is an alternative implementation of saturation adjustment, attempting to -bypass compilation limitations of the gtfn backend. This implementation is less -elegant than the main implementation in thermo.py. Neither is currently used -in the test code run_graupel_only -- the only implementation which has passed -validation -- and the run_full_muphys (currently untested) uses the standard -thermo.py implementation, in the hope that the gtfn compilation problems will -disappear soon. -""" - -import gt4py.next as gtx -from gt4py.next import maximum, where - -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.frozen import g_ct, t_d -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import ( - _dqsatdT_rho, - _qsat_rho, -) -from icon4py.model.common import field_type_aliases as fa, type_alias as ta - - -@gtx.field_operator -def _satadj_init( - te: fa.CellKField[ta.wpfloat], # Temperature - qve: fa.CellKField[ta.wpfloat], # Specific humidity - qce: fa.CellKField[ta.wpfloat], # Specific cloud water content - qre: fa.CellKField[ta.wpfloat], # Specific rain water - qti: fa.CellKField[ta.wpfloat], # Specific mass of all ice species (total-ice) -) -> tuple[ - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], -]: - qt = qve + qce + qre + qti # temporary, used only here - cvc = t_d.cvd * (1.0 - qt) + t_d.clw * qre + g_ct.ci * qti # output variable - cv = cvc + t_d.cvv * qve + t_d.clw * qce # temporary, used only here - ue = cv * te - qce * g_ct.lvc # output variable - Tx_hold = ue / (cv + qce * (t_d.cvv - t_d.clw)) - Tx = te - return cvc, ue, Tx_hold, Tx # output variables - - -@gtx.field_operator -def _output_calculation( - qve: fa.CellKField[ta.wpfloat], # Specific humidity - qce: fa.CellKField[ta.wpfloat], # Specific cloud water content - qx_hold: fa.CellKField[ta.wpfloat], # TBD - qx: fa.CellKField[ta.wpfloat], # TBD - Tx_hold: fa.CellKField[ta.wpfloat], # TBD - Tx: fa.CellKField[ta.wpfloat], # TBD -) -> tuple[ - fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat] -]: # Internal energy - te = where((qve + qce <= qx_hold), Tx_hold, Tx) - qce = where((qve + qce <= qx_hold), 0.0, maximum(qve + qce - qx, 0.0)) - qve = where((qve + qce <= qx_hold), qve + qce, qx) - return te, qve, qce - - -@gtx.field_operator -def _newton_raphson( - qx: fa.CellKField[ta.wpfloat], - dqx: fa.CellKField[ta.wpfloat], - Tx: fa.CellKField[ta.wpfloat], - rho: fa.CellKField[ta.wpfloat], - qve: fa.CellKField[ta.wpfloat], - qce: fa.CellKField[ta.wpfloat], - cvc: fa.CellKField[ta.wpfloat], - ue: fa.CellKField[ta.wpfloat], -) -> fa.CellKField[ta.wpfloat]: - qcx = qve + qce - qx - cv = cvc + t_d.cvv * qx + t_d.clw * qcx - ux = cv * Tx - qcx * g_ct.lvc - dux = cv + dqx * (g_ct.lvc + (t_d.cvv - t_d.clw) * Tx) - Tx = Tx - (ux - ue) / dux - return Tx - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def saturation_adjustment2( - te: fa.CellKField[ta.wpfloat], # Temperature - qve: fa.CellKField[ta.wpfloat], # Specific humidity - qce: fa.CellKField[ta.wpfloat], # Specific cloud water content - qre: fa.CellKField[ta.wpfloat], # Specific rain water - qti: fa.CellKField[ta.wpfloat], # Specific mass of all ice species (total-ice) - rho: fa.CellKField[ta.wpfloat], # Density containing dry air and water constituents - cvc: fa.CellKField[ta.wpfloat], # Temporary field - ue: fa.CellKField[ta.wpfloat], # Temporary field - Tx_hold: fa.CellKField[ta.wpfloat], # Temporary field - Tx: fa.CellKField[ta.wpfloat], # Temporary field - qx_hold: fa.CellKField[ta.wpfloat], # Temporary field - qx: fa.CellKField[ta.wpfloat], # Temporary field - dqx: fa.CellKField[ta.wpfloat], # Temporary field - qve_out: fa.CellKField[ta.wpfloat], # Specific humidity - qce_out: fa.CellKField[ta.wpfloat], # Specific cloud water content - te_out: fa.CellKField[ta.wpfloat], # Temperature -): - _satadj_init(te, qve, qce, qre, qti, out=(cvc, ue, Tx_hold, Tx)) - _qsat_rho(Tx_hold, rho, out=qx_hold) - - # Newton-Raphson iteration - _qsat_rho(Tx, rho, out=qx) - _dqsatdT_rho(qx, Tx, out=dqx) - _newton_raphson(qx, dqx, Tx, rho, qve, qce, cvc, ue, out=Tx) - _qsat_rho(Tx, rho, out=qx) - _dqsatdT_rho(qx, Tx, out=dqx) - _newton_raphson(qx, dqx, Tx, rho, qve, qce, cvc, ue, out=Tx) - _qsat_rho(Tx, rho, out=qx) - _dqsatdT_rho(qx, Tx, out=dqx) - _newton_raphson(qx, dqx, Tx, rho, qve, qce, cvc, ue, out=Tx) - _qsat_rho(Tx, rho, out=qx) - _dqsatdT_rho(qx, Tx, out=dqx) - _newton_raphson(qx, dqx, Tx, rho, qve, qce, cvc, ue, out=Tx) - _qsat_rho(Tx, rho, out=qx) - _dqsatdT_rho(qx, Tx, out=dqx) - _newton_raphson(qx, dqx, Tx, rho, qve, qce, cvc, ue, out=Tx) - _qsat_rho(Tx, rho, out=qx) - _dqsatdT_rho(qx, Tx, out=dqx) - _newton_raphson(qx, dqx, Tx, rho, qve, qce, cvc, ue, out=Tx) - - # final humidity calculation - _qsat_rho(Tx, rho, out=qx) - - # final calculation of output variables - _output_calculation(qve, qce, qx_hold, qx, Tx_hold, Tx, out=(te_out, qve_out, qce_out)) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py index 31242b015c..3e4a479416 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/core/thermo.py @@ -6,10 +6,10 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause import gt4py.next as gtx -from gt4py.next import exp, maximum, where +from gt4py.next import exp from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.frozen import g_ct, t_d -from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta +from icon4py.model.common import field_type_aliases as fa, type_alias as ta @gtx.field_operator @@ -338,99 +338,3 @@ def _newton_raphson( dux = cv + dqx * (g_ct.lvc + (t_d.cvv - t_d.clw) * Tx) Tx = Tx - (ux - ue) / dux return Tx - - -@gtx.field_operator -def _saturation_adjustment( - te: fa.CellKField[ta.wpfloat], - qve: fa.CellKField[ta.wpfloat], - qce: fa.CellKField[ta.wpfloat], - qre: fa.CellKField[ta.wpfloat], - qse: fa.CellKField[ta.wpfloat], - qie: fa.CellKField[ta.wpfloat], - qge: fa.CellKField[ta.wpfloat], - rho: fa.CellKField[ta.wpfloat], -) -> tuple[ - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], -]: - """ - Compute the saturation adjustment which revises internal energy and water contents - - Args: - Tx: Temperature - qve: Specific humidity - qce: Specific cloud water content - qre: Specific rain water - qti: Specific mass of all ice species (total-ice) - rho: Density containing dry air and water constituents - - Result: Tuple containing - - Revised temperature - - Revised specific cloud water content - - Revised specific vapor content - - Mask specifying where qce+qve less than holding capacity - """ - qti = qse + qie + qge - qt = qve + qce + qre + qti - cvc = t_d.cvd * (1.0 - qt) + t_d.clw * qre + g_ct.ci * qti - cv = cvc + t_d.cvv * qve + t_d.clw * qce - ue = cv * te - qce * g_ct.lvc - Tx_hold = ue / (cv + qce * (t_d.cvv - t_d.clw)) - qx_hold = _qsat_rho(Tx_hold, rho) - - Tx = te - # Newton-Raphson iteration: 6 times the same operations - Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue) - Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue) - Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue) - Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue) - Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue) - Tx = _newton_raphson(Tx, rho, qve, qce, cvc, ue) - - # At this point we hope Tx has converged - qx = _qsat_rho(Tx, rho) - - # Is it possible to unify the where for all three outputs?? - mask = qve + qce <= qx_hold - te = where(mask, Tx_hold, Tx) - qce = where(mask, 0.0, maximum(qve + qce - qx, 0.0)) - qve = where(mask, qve + qce, qx) - - return te, qve, qce - - -@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) -def saturation_adjustment( - te: fa.CellKField[ta.wpfloat], # Temperature - qve: fa.CellKField[ta.wpfloat], # Specific humidity - qce: fa.CellKField[ta.wpfloat], # Specific cloud water content - qre: fa.CellKField[ta.wpfloat], # Specific rain water - qse: fa.CellKField[ta.wpfloat], # Specific snow water - qie: fa.CellKField[ta.wpfloat], # Specific ice water content - qge: fa.CellKField[ta.wpfloat], # Specific graupel water content - rho: fa.CellKField[ta.wpfloat], # Density containing dry air and water constituents - te_out: fa.CellKField[ta.wpfloat], # Temperature - qve_out: fa.CellKField[ta.wpfloat], # Specific humidity - qce_out: fa.CellKField[ta.wpfloat], # Specific cloud water content - horizontal_start: gtx.int32, - horizontal_end: gtx.int32, - vertical_start: gtx.int32, - vertical_end: gtx.int32, -): - _saturation_adjustment( - te, - qve, - qce, - qre, - qse, - qie, - qge, - rho, - out=(te_out, qve_out, qce_out), - domain={ - dims.CellDim: (horizontal_start, horizontal_end), - dims.KDim: (vertical_start, vertical_end), - }, - ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/common.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/common.py new file mode 100644 index 0000000000..f5ee93f063 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/common.py @@ -0,0 +1,216 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import dataclasses +import functools +import pathlib + +import netCDF4 +import numpy as np +from gt4py import next as gtx +from gt4py.next import typing as gtx_typing + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import Q +from icon4py.model.common import dimension as dims + + +def _calc_dz(z: np.ndarray) -> np.ndarray: + ksize = z.shape[0] + dz = np.zeros(z.shape, np.float64) + zh = 1.5 * z[ksize - 1, :] - 0.5 * z[ksize - 2, :] + for k in range(ksize - 1, -1, -1): + zh_new = 2.0 * z[k, :] - zh + dz[k, :] = -zh + zh_new + zh = zh_new + return dz + + +def _as_field_from_nc( + dataset: netCDF4.Dataset, + allocator: gtx_typing.FieldBufferAllocationUtil, + varname: str, + optional: bool = False, + dtype: np.dtype | None = None, +) -> gtx.Field[dims.CellDim, dims.KDim] | None: + if optional and varname not in dataset.variables: + return None + + var = dataset.variables[varname] + if var.dimensions[0] == "time": + var = var[0, :, :] + data = np.transpose(var) + if dtype is not None: + data = data.astype(dtype) + return gtx.as_field( + (dims.CellDim, dims.KDim), + data, + allocator=allocator, + ) + + +def _field_to_nc( + dataset: netCDF4.Dataset, + dims: tuple[str, str], + varname: str, + field: gtx.Field[dims.CellDim, dims.KDim], + dtype: np.dtype = np.float64, +) -> None: + var = dataset.createVariable(varname, dtype, dims) + var[...] = field.asnumpy().transpose() + + +@dataclasses.dataclass +class GraupelInput: + ncells: int + nlev: int + dz: gtx.Field[dims.CellDim, dims.KDim] + p: gtx.Field[dims.CellDim, dims.KDim] + rho: gtx.Field[dims.CellDim, dims.KDim] + t: gtx.Field[dims.CellDim, dims.KDim] + qv: gtx.Field[dims.CellDim, dims.KDim] + qc: gtx.Field[dims.CellDim, dims.KDim] + qi: gtx.Field[dims.CellDim, dims.KDim] + qr: gtx.Field[dims.CellDim, dims.KDim] + qs: gtx.Field[dims.CellDim, dims.KDim] + qg: gtx.Field[dims.CellDim, dims.KDim] + + @property + def q(self) -> Q: + return Q( + v=self.qv, + c=self.qc, + r=self.qr, + s=self.qs, + i=self.qi, + g=self.qg, + ) + + @classmethod + def load( + cls, filename: pathlib.Path | str, allocator: gtx_typing.FieldBufferAllocationUtil + ) -> None: + with netCDF4.Dataset(filename, mode="r") as ncfile: + try: + ncells = len(ncfile.dimensions["cell"]) + except KeyError: + ncells = len(ncfile.dimensions["ncells"]) + + nlev = len(ncfile.dimensions["height"]) + + dz = _calc_dz(ncfile.variables["zg"]) + + field_from_nc = functools.partial( + _as_field_from_nc, ncfile, allocator, dtype=np.float64 + ) + return cls( + ncells=ncells, + nlev=nlev, + dz=gtx.as_field((dims.CellDim, dims.KDim), np.transpose(dz), allocator=allocator), + t=field_from_nc("ta"), + p=field_from_nc("pfull"), + qs=field_from_nc("qs"), + qi=field_from_nc("cli"), + qg=field_from_nc("qg"), + qv=field_from_nc("hus"), + qc=field_from_nc("clw"), + qr=field_from_nc("qr"), + rho=field_from_nc("rho"), + ) + + +@dataclasses.dataclass +class GraupelOutput: + t: gtx.Field[dims.CellDim, dims.KDim] + qv: gtx.Field[dims.CellDim, dims.KDim] + qc: gtx.Field[dims.CellDim, dims.KDim] + qi: gtx.Field[dims.CellDim, dims.KDim] + qr: gtx.Field[dims.CellDim, dims.KDim] + qs: gtx.Field[dims.CellDim, dims.KDim] + qg: gtx.Field[dims.CellDim, dims.KDim] + + pflx: gtx.Field[dims.CellDim, dims.KDim] | None + pr: gtx.Field[dims.CellDim, dims.KDim] | None + ps: gtx.Field[dims.CellDim, dims.KDim] | None + pi: gtx.Field[dims.CellDim, dims.KDim] | None + pg: gtx.Field[dims.CellDim, dims.KDim] | None + pre: gtx.Field[dims.CellDim, dims.KDim] | None + + @classmethod + def allocate(cls, allocator: gtx_typing.FieldBufferAllocationUtil, domain: gtx.Domain): + zeros = functools.partial(gtx.zeros, domain=domain, allocator=allocator) + # TODO +1 size fields? + return cls(**{field.name: zeros() for field in dataclasses.fields(cls)}) + + @classmethod + def load(cls, filename: pathlib.Path | str, allocator: gtx_typing.FieldBufferAllocationUtil): + with netCDF4.Dataset(filename, mode="r") as ncfile: + field_from_nc = functools.partial(_as_field_from_nc, ncfile, allocator) + return cls( + t=field_from_nc("ta"), + qv=field_from_nc("hus"), + qc=field_from_nc("clw"), + qi=field_from_nc("cli"), + qr=field_from_nc("qr"), + qs=field_from_nc("qs"), + qg=field_from_nc("qg"), + pflx=field_from_nc("pflx", optional=True), + pr=field_from_nc("prr_gsp", optional=True), + ps=field_from_nc("prs_gsp", optional=True), + pi=field_from_nc("pri_gsp", optional=True), + pg=field_from_nc("prg_gsp", optional=True), + pre=field_from_nc("pre_gsp", optional=True), + ) + + @property + def q(self) -> Q: + return Q( + v=self.qv, + c=self.qc, + r=self.qr, + s=self.qs, + i=self.qi, + g=self.qg, + ) + + def write(self, filename: pathlib.Path | str): + ncells = self.t.shape[0] + nlev = self.t.shape[1] + + with netCDF4.Dataset(filename, mode="w") as ncfile: + ncfile.createDimension("ncells", ncells) + ncfile.createDimension("height", nlev) + ncfile.createDimension("height1", nlev + 1) # what's the reason for the +1 fields here? + + write_height_field = functools.partial( + _field_to_nc, ncfile, ("height", "ncells"), dtype=np.float64 + ) + write_height1_field = functools.partial( # TODO + _field_to_nc, ncfile, ("height1", "ncells"), dtype=np.float64 + ) + + write_height_field("ta", self.t) + write_height_field("hus", self.qv) + write_height_field("clw", self.qc) + write_height_field("cli", self.qi) + write_height_field("qr", self.qr) + write_height_field("qs", self.qs) + write_height_field("qg", self.qg) + if self.pflx is not None: + write_height_field("pflx", self.pflx) + if self.pr is not None: + write_height_field("prr_gsp", self.pr) # TODO height1? + if self.ps is not None: + write_height_field("prs_gsp", self.ps) # TODO + if self.pi is not None: + write_height_field("pri_gsp", self.pi) # TODO + if self.pg is not None: + write_height_field("prg_gsp", self.pg) # TODO + if self.pre is not None: + write_height_field("pre_gsp", self.pre) # TODO diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py index c9d57b3b60..f67446d810 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py @@ -7,511 +7,211 @@ # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause -""" -This is the full muphys implementation for a single muphys call -WORK IN PROGRESS!!!! Do not try to run this. -""" - -# TODO refactor similar to run_graupel_only +from __future__ import annotations import argparse -import sys +import functools +import pathlib import time - -import gt4py.next as gtx -import netCDF4 -import numpy as np - -from icon4py.model.common import model_options -from icon4py.model.common.model_options import setup_program - - -try: - from netCDF4 import Dataset -except ImportError: - print("Netcdf not installed") - sys.exit() - -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import saturation_adjustment -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations.graupel import ( - graupel_run, +from collections.abc import Callable + +from gt4py import next as gtx + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core import saturation_adjustment +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import common, utils +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations import graupel, muphys +from icon4py.model.common import ( + dimension as dims, + field_type_aliases as fa, + model_backends, + model_options, + type_alias as ta, ) -from icon4py.model.common import dimension as dims, model_backends +from icon4py.model.common.utils import device_utils -def set_lib_path(lib_dir): - sys.path.append(lib_dir) +# TODO(havogt): make similar to icon4py driver structure def get_args(): parser = argparse.ArgumentParser() parser.add_argument( - "-o", - metavar="output_file", - dest="output_file", - help="output filename", - default="output.nc", + "-o", metavar="output_file", dest="output_file", help="output filename", default="output.nc" ) parser.add_argument( - "-b", - metavar="backend", - dest="backend", - help="gt4py backend", - default="gtfn_cpu", + "-b", metavar="backend", dest="backend", help="gt4py backend", default="gtfn_cpu" ) parser.add_argument("input_file", help="input data file") parser.add_argument("itime", help="time-index", nargs="?", default=0) parser.add_argument("dt", help="timestep", nargs="?", default=30.0) parser.add_argument("qnc", help="Water number concentration", nargs="?", default=100.0) - parser.add_argument( - "-ldir", - metavar="lib_dir", - dest="ldir", - help="directory with py_graupel shared lib", - default="build/lib64", - ) - parser.add_argument( - "-with_sat_adj", - action="store_true", - ) return parser.parse_args() -class Data: - def __init__(self, args): - nc = netCDF4.Dataset(args.input_file) - # intent(in) variables: - try: - self.ncells = len(nc.dimensions["cell"]) - except KeyError: - self.ncells = len(nc.dimensions["ncells"]) - - self.nlev = len(nc.dimensions["height"]) - self.z = nc.variables["zg"][:, :].astype(np.float64) - self.p = nc.variables["pfull"][:, :].astype(np.float64) - self.rho = nc.variables["rho"][:, :].astype(np.float64) - # intent(inout) variables: - self.t = nc.variables["ta"][:, :].astype(np.float64) # inout - self.qv = nc.variables["hus"][:, :].astype(np.float64) # inout - self.qc = nc.variables["clw"][:, :].astype(np.float64) # inout - self.qi = nc.variables["cli"][:, :].astype(np.float64) # inout - self.qr = nc.variables["qr"][:, :].astype(np.float64) # inout - self.qs = nc.variables["qs"][:, :].astype(np.float64) # inout - self.qg = nc.variables["qg"][:, :].astype(np.float64) # inout - # intent(out) variables: - self.t_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qv_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qc_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qi_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qr_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qs_out = np.zeros((self.ncells, self.nlev), np.float64) - self.qg_out = np.zeros((self.ncells, self.nlev), np.float64) - self.pflx_out = np.zeros((self.ncells, self.nlev), np.float64) - self.prr_gsp = np.zeros(self.ncells, np.float64) - self.pri_gsp = np.zeros(self.ncells, np.float64) - self.prs_gsp = np.zeros(self.ncells, np.float64) - self.prg_gsp = np.zeros(self.ncells, np.float64) - self.pre_gsp = np.zeros(self.ncells, np.float64) - self.dz = calc_dz(self.nlev, self.z) - self.mask_out = np.full((self.ncells, self.nlev), True) - - -def calc_dz(ksize, z): - dz = np.zeros(z.shape, np.float64) - zh = 1.5 * z[ksize - 1, :] - 0.5 * z[ksize - 2, :] - for k in range(ksize - 1, -1, -1): - zh_new = 2.0 * z[k, :] - zh - dz[k, :] = -zh + zh_new - zh = zh_new - return dz - - -def run_program( - args, backend, data, t_out, qv_out, qc_out, qi_out, qr_out, qs_out, qg_out, pflx_out +def _muphys_step_separate( + graupel_program: Callable, + saturation_adjustment_program: Callable, + dz: fa.CellKField[ta.wpfloat], + te: fa.CellKField[ta.wpfloat], # Temperature + p: fa.CellKField[ta.wpfloat], # Pressure + rho: fa.CellKField[ta.wpfloat], # Density containing dry air and water constituents + q_in: common.Q, + q_out: common.Q, + t_out: fa.CellKField[ta.wpfloat], # Revised temperature + pflx: fa.CellKField[ta.wpfloat], # Total precipitation flux + pr: fa.CellKField[ta.wpfloat], # Precipitation of rain + ps: fa.CellKField[ta.wpfloat], # Precipitation of snow + pi: fa.CellKField[ta.wpfloat], # Precipitation of ice + pg: fa.CellKField[ta.wpfloat], # Precipitation of graupel + pre: fa.CellKField[ta.wpfloat], # Precipitation of graupel ): - ksize = data.dz.shape[0] - - dz = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.dz[:, :]), - allocator=backend, - ) - te = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.t[0, :, :]), - allocator=backend, - ) - p = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.p[0, :, :]), - allocator=backend, - ) - qse = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qs[0, :, :]), - allocator=backend, - ) - qie = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qi[0, :, :]), - allocator=backend, - ) - qge = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qg[0, :, :]), - allocator=backend, - ) - qve = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qv[0, :, :]), - allocator=backend, - ) - qce = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qc[0, :, :]), - allocator=backend, - ) - qre = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.qr[0, :, :]), - allocator=backend, - ) - rho = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.transpose(data.rho[0, :, :]), - allocator=backend, - ) - - pr_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, - ) - ps_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, - ) - pi_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, - ) - pg_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, - ) - pre_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - np.zeros((data.ncells, data.nlev)), - allocator=backend, - ) - mask_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.mask_out, - allocator=backend, - ) - - saturation_adjustment_program = setup_program( - backend=backend, - program=saturation_adjustment, - constant_args={}, - horizontal_sizes={ - "horizontal_start": gtx.int32(0), - "horizontal_end": data.ncells, - }, - vertical_sizes={ - "vertical_start": gtx.int32(0), - "vertical_end": gtx.int32(data.nlev), - }, - offset_provider={"Koff": dims.KDim}, - ) - - graupel_run_program = setup_program( - backend=backend, - program=graupel_run, - constant_args={}, - horizontal_sizes={ - "horizontal_start": gtx.int32(0), - "horizontal_end": data.ncells, - }, - vertical_sizes={ - "vertical_start": gtx.int32(0), - "vertical_end": gtx.int32(data.nlev), - "last_lev": ksize - 1, - }, - offset_provider={"Koff": dims.KDim}, - ) - - start_time = time.time() - + # In-place update ok since saturation_adjustment is fully point-wise, + # but not recommended. TODO + saturation_adjustment_program( + te=te, + q_in=q_in, + rho=rho, + te_out=te, + qve_out=q_in.v, + qce_out=q_in.c, + ) + + graupel_program( + dz=dz, + te=te, + p=p, + rho=rho, + q_in=q_in, + t_out=t_out, + q_out=q_out, + pflx=pflx, + pr=pr, + ps=ps, + pi=pi, + pg=pg, + pre=pre, + ) + + saturation_adjustment_program( + te=t_out, + q_in=q_out, + rho=rho, + te_out=t_out, + qve_out=q_out.v, + qce_out=q_out.c, + ) + + +def setup_muphys( + inp: common.GraupelInput, + dt: float, + qnc: float, + backend: model_backends.BackendLike, + *, + single_program: bool = False, +): + if single_program: + # TODO(havogt): make an option in gt4py for thread-safety? + with utils.recursion_limit(10**5): + muphys_program = model_options.setup_program( + backend=backend, + program=muphys.muphys_run, + constant_args={"dt": dt, "qnc": qnc}, + horizontal_sizes={ + "horizontal_start": gtx.int32(0), + "horizontal_end": inp.ncells, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(inp.nlev), + }, + offset_provider={"Koff": dims.KDim}, + ) + gtx.wait_for_compilation() + return muphys_program + else: + with utils.recursion_limit(10**5): # TODO(havogt): make an option in gt4py? + graupel_run_program = model_options.setup_program( + backend=backend, + program=graupel.graupel_run, + constant_args={"dt": dt, "qnc": qnc}, + horizontal_sizes={ + "horizontal_start": gtx.int32(0), + "horizontal_end": inp.ncells, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(inp.nlev), + }, + offset_provider={"Koff": dims.KDim}, + ) + saturation_adjustment_program = model_options.setup_program( + backend=backend, + program=saturation_adjustment.saturation_adjustment, + horizontal_sizes={ + "horizontal_start": gtx.int32(0), + "horizontal_end": inp.ncells, + }, + vertical_sizes={ + "vertical_start": gtx.int32(0), + "vertical_end": gtx.int32(inp.nlev), + }, + ) + gtx.wait_for_compilation() + + return functools.partial( + _muphys_step_separate, + graupel_program=graupel_run_program, + saturation_adjustment_program=saturation_adjustment_program, + ) + + +def main(): + args = get_args() + + backend = model_backends.BACKENDS[args.backend] + allocator = model_backends.get_allocator(backend) + + inp = common.GraupelInput.load(filename=pathlib.Path(args.input_file), allocator=allocator) + out = common.GraupelOutput.allocate( + domain=gtx.domain({dims.CellDim: inp.ncells, dims.KDim: inp.nlev}), allocator=allocator + ) + + # TODO(havogt): once we see single program being equally fast, remove the other implementation + muphys_step = setup_muphys(inp, dt=args.dt, qnc=args.qnc, backend=backend, single_program=False) + + start_time = None for _x in range(int(args.itime) + 1): if _x == 1: # Only start timing second iteration + device_utils.sync(backend) start_time = time.time() - saturation_adjustment_program( - te=te, - qve=qve, - qce=qce, - qre=qre, - qse=qse, - qie=qie, - qge=qge, - rho=rho, - te_out=t_out, # Temperature - qve_out=qv_out, # Specific humidity - qce_out=qc_out, # Specific cloud water content - mask_out=mask_out, # Mask of interest - ) - - graupel_run_program( - dz=dz, - te=te, - p=p, - rho=rho, - qve=qve, - qce=qce, - qre=qre, - qse=qse, - qie=qie, - qge=qge, - dt=args.dt, - qnc=args.qnc, - t_out=t_out, - qv_out=qv_out, - qc_out=qc_out, - qr_out=qr_out, - qs_out=qs_out, - qi_out=qi_out, - qg_out=qg_out, - pflx=pflx_out, - pr=pr_out, - ps=ps_out, - pi=pi_out, - pg=pg_out, - pre=pre_out, + muphys_step( + dz=inp.dz, + te=inp.t, + p=inp.p, + rho=inp.rho, + q_in=inp.q, + q_out=out.q, + t_out=out.t, + pflx=out.pflx, + pr=out.pr, + ps=out.ps, + pi=out.pi, + pg=out.pg, + pre=out.pre, ) - saturation_adjustment_program( - te=te, - qve=qve, - qce=qce, - qre=qre, - qse=qse, - qie=qie, - qge=qge, - rho=rho, - te_out=t_out, # Temperature - qve_out=qv_out, # Specific humidity - qce_out=qc_out, # Specific cloud water content - mask_out=mask_out, # Mask of interest - ) - - if _x == int(args.itime): # End timer on last iteration - end_time = time.time() - - elapsed_time = end_time - start_time - print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") - data.prr_gsp = np.transpose(pr_out[dims.KDim(ksize - 1)].asnumpy()) - data.prs_gsp = np.transpose(ps_out[dims.KDim(ksize - 1)].asnumpy()) - data.pri_gsp = np.transpose(pi_out[dims.KDim(ksize - 1)].asnumpy()) - data.prg_gsp = np.transpose(pg_out[dims.KDim(ksize - 1)].asnumpy()) - data.pre_gsp = np.transpose(pre_out[dims.KDim(ksize - 1)].asnumpy()) - + device_utils.sync(backend) + end_time = time.time() -def write_fields( - output_filename, - ncell, - nlev, - t, - qv, - qc, - qi, - qr, - qs, - qg, - prr_gsp, - prs_gsp, - pri_gsp, - prg_gsp, - pflx, - pre_gsp, -): - ncfile = Dataset(output_filename, mode="w") - ncells = ncfile.createDimension("ncells", ncell) - height = ncfile.createDimension("height", nlev) - height1 = ncfile.createDimension("height1", nlev + 1) - ta_var = ncfile.createVariable("ta", np.double, ("height", "ncells")) - hus_var = ncfile.createVariable("hus", np.double, ("height", "ncells")) - clw_var = ncfile.createVariable("clw", np.double, ("height", "ncells")) - cli_var = ncfile.createVariable("cli", np.double, ("height", "ncells")) - qr_var = ncfile.createVariable("qr", np.double, ("height", "ncells")) - qs_var = ncfile.createVariable("qs", np.double, ("height", "ncells")) - qg_var = ncfile.createVariable("qg", np.double, ("height", "ncells")) - pflx_var = ncfile.createVariable("pflx", np.double, ("height", "ncells")) - prr_gsp_var = ncfile.createVariable("prr_gsp", np.double, ("height1", "ncells")) - prs_gsp_var = ncfile.createVariable("prs_gsp", np.double, ("height1", "ncells")) - pri_gsp_var = ncfile.createVariable("pri_gsp", np.double, ("height1", "ncells")) - prg_gsp_var = ncfile.createVariable("prg_gsp", np.double, ("height1", "ncells")) - pre_gsp_var = ncfile.createVariable("pre_gsp", np.double, ("height1", "ncells")) - - ta_var[:, :] = t - hus_var[:, :] = qv - clw_var[:, :] = qc - cli_var[:, :] = qi - qr_var[:, :] = qr - qs_var[:, :] = qs - qg_var[:, :] = qg - pflx_var[:, :] = pflx - prr_gsp_var[:, :] = prr_gsp - prs_gsp_var[:, :] = prs_gsp - pri_gsp_var[:, :] = pri_gsp - prg_gsp_var[:, :] = prg_gsp - pre_gsp_var[:, :] = pre_gsp - ncfile.close() - - -args = get_args() -backend_like = model_backends.BACKENDS[args.backend] -backend = model_options.customize_backend(None, backend_like) + if start_time is not None: + elapsed_time = end_time - start_time + print("For", int(args.itime), "iterations it took", elapsed_time, "seconds!") -set_lib_path(args.ldir) -sys.setrecursionlimit(10**4) + out.write(args.output_file) -data = Data(args) -t_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.t_out, - allocator=backend, -) -qv_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qv_out, - allocator=backend, -) -qc_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qc_out, - allocator=backend, -) -qr_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qr_out, - allocator=backend, -) -qs_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qs_out, - allocator=backend, -) -qi_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qi_out, - allocator=backend, -) -qg_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.qg_out, - allocator=backend, -) -pflx_out = gtx.as_field( - ( - dims.CellDim, - dims.KDim, - ), - data.pflx_out, - allocator=backend, -) - -run_program(args, backend, data, t_out, qv_out, qc_out, qi_out, qr_out, qs_out, qg_out, pflx_out) - -write_fields( - args.output_file, - data.ncells, - data.nlev, - t=np.transpose(t_out.asnumpy()), - qv=np.transpose(qv_out.asnumpy()), - qc=np.transpose(qc_out.asnumpy()), - qi=np.transpose(qi_out.asnumpy()), - qr=np.transpose(qr_out.asnumpy()), - qs=np.transpose(qs_out.asnumpy()), - qg=np.transpose(qg_out.asnumpy()), - prr_gsp=data.prr_gsp, - pri_gsp=data.pri_gsp, - prs_gsp=data.prs_gsp, - prg_gsp=data.prg_gsp, - pflx=np.transpose(pflx_out.asnumpy()), - pre_gsp=data.pre_gsp, -) +if __name__ == "__main__": + main() diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py index 1c67c86c21..865397d86b 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py @@ -10,17 +10,12 @@ from __future__ import annotations import argparse -import dataclasses -import functools import pathlib import time -import netCDF4 -import numpy as np from gt4py import next as gtx -from gt4py.next import typing as gtx_typing -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import utils +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import common, utils from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations import graupel from icon4py.model.common import dimension as dims, model_backends, model_options from icon4py.model.common.utils import device_utils @@ -45,180 +40,10 @@ def get_args(): return parser.parse_args() -# TODO double check the sizes of pxxx vars (should be nlev+1?) - - -def _calc_dz(z: np.ndarray) -> np.ndarray: - ksize = z.shape[0] - dz = np.zeros(z.shape, np.float64) - zh = 1.5 * z[ksize - 1, :] - 0.5 * z[ksize - 2, :] - for k in range(ksize - 1, -1, -1): - zh_new = 2.0 * z[k, :] - zh - dz[k, :] = -zh + zh_new - zh = zh_new - return dz - - -def _as_field_from_nc( - dataset: netCDF4.Dataset, - allocator: gtx_typing.FieldBufferAllocationUtil, - varname: str, - optional: bool = False, - dtype: np.dtype | None = None, -) -> gtx.Field[dims.CellDim, dims.KDim] | None: - if optional and varname not in dataset.variables: - return None - - var = dataset.variables[varname] - if var.dimensions[0] == "time": - var = var[0, :, :] - data = np.transpose(var) - if dtype is not None: - data = data.astype(dtype) - return gtx.as_field( - (dims.CellDim, dims.KDim), - data, - allocator=allocator, - ) - - -def _field_to_nc( - dataset: netCDF4.Dataset, - dims: tuple[str, str], - varname: str, - field: gtx.Field[dims.CellDim, dims.KDim], - dtype: np.dtype = np.float64, -) -> None: - var = dataset.createVariable(varname, dtype, dims) - var[...] = field.asnumpy().transpose() - - -@dataclasses.dataclass -class GraupelInput: - ncells: int - nlev: int - dz: gtx.Field[dims.CellDim, dims.KDim] - p: gtx.Field[dims.CellDim, dims.KDim] - rho: gtx.Field[dims.CellDim, dims.KDim] - t: gtx.Field[dims.CellDim, dims.KDim] - qv: gtx.Field[dims.CellDim, dims.KDim] - qc: gtx.Field[dims.CellDim, dims.KDim] - qi: gtx.Field[dims.CellDim, dims.KDim] - qr: gtx.Field[dims.CellDim, dims.KDim] - qs: gtx.Field[dims.CellDim, dims.KDim] - qg: gtx.Field[dims.CellDim, dims.KDim] - - @classmethod - def load( - cls, filename: pathlib.Path | str, allocator: gtx_typing.FieldBufferAllocationUtil - ) -> None: - with netCDF4.Dataset(filename, mode="r") as ncfile: - try: - ncells = len(ncfile.dimensions["cell"]) - except KeyError: - ncells = len(ncfile.dimensions["ncells"]) - - nlev = len(ncfile.dimensions["height"]) - - dz = _calc_dz(ncfile.variables["zg"]) - - field_from_nc = functools.partial( - _as_field_from_nc, ncfile, allocator, dtype=np.float64 - ) - return cls( - ncells=ncells, - nlev=nlev, - dz=gtx.as_field((dims.CellDim, dims.KDim), np.transpose(dz), allocator=allocator), - t=field_from_nc("ta"), - p=field_from_nc("pfull"), - qs=field_from_nc("qs"), - qi=field_from_nc("cli"), - qg=field_from_nc("qg"), - qv=field_from_nc("hus"), - qc=field_from_nc("clw"), - qr=field_from_nc("qr"), - rho=field_from_nc("rho"), - ) - - -@dataclasses.dataclass -class GraupelOutput: - t: gtx.Field[dims.CellDim, dims.KDim] - qv: gtx.Field[dims.CellDim, dims.KDim] - qc: gtx.Field[dims.CellDim, dims.KDim] - qi: gtx.Field[dims.CellDim, dims.KDim] - qr: gtx.Field[dims.CellDim, dims.KDim] - qs: gtx.Field[dims.CellDim, dims.KDim] - qg: gtx.Field[dims.CellDim, dims.KDim] - - pflx: gtx.Field[dims.CellDim, dims.KDim] | None - pr: gtx.Field[dims.CellDim, dims.KDim] | None - ps: gtx.Field[dims.CellDim, dims.KDim] | None - pi: gtx.Field[dims.CellDim, dims.KDim] | None - pg: gtx.Field[dims.CellDim, dims.KDim] | None - pre: gtx.Field[dims.CellDim, dims.KDim] | None - - @classmethod - def allocate(cls, allocator: gtx_typing.FieldBufferAllocationUtil, domain: gtx.Domain): - zeros = functools.partial(gtx.zeros, domain=domain, allocator=allocator) - # TODO +1 size fields? - return cls(**{field.name: zeros() for field in dataclasses.fields(cls)}) - - @classmethod - def load(cls, filename: pathlib.Path | str, allocator: gtx_typing.FieldBufferAllocationUtil): - with netCDF4.Dataset(filename, mode="r") as ncfile: - field_from_nc = functools.partial(_as_field_from_nc, ncfile, allocator) - return cls( - t=field_from_nc("ta"), - qv=field_from_nc("hus"), - qc=field_from_nc("clw"), - qi=field_from_nc("cli"), - qr=field_from_nc("qr"), - qs=field_from_nc("qs"), - qg=field_from_nc("qg"), - pflx=field_from_nc("pflx", optional=True), - pr=field_from_nc("prr_gsp", optional=True), - ps=field_from_nc("prs_gsp", optional=True), - pi=field_from_nc("pri_gsp", optional=True), - pg=field_from_nc("prg_gsp", optional=True), - pre=field_from_nc("pre_gsp", optional=True), - ) - - def write(self, filename: pathlib.Path | str): - ncells = self.t.shape[0] - nlev = self.t.shape[1] - - with netCDF4.Dataset(filename, mode="w") as ncfile: - ncfile.createDimension("ncells", ncells) - ncfile.createDimension("height", nlev) - - write_height_field = functools.partial( - _field_to_nc, ncfile, ("height", "ncells"), dtype=np.float64 - ) - - write_height_field("ta", self.t) - write_height_field("hus", self.qv) - write_height_field("clw", self.qc) - write_height_field("cli", self.qi) - write_height_field("qr", self.qr) - write_height_field("qs", self.qs) - write_height_field("qg", self.qg) - if self.pflx is not None: - write_height_field("pflx", self.pflx) - if self.pr is not None: - write_height_field("prr_gsp", self.pr) - if self.ps is not None: - write_height_field("prs_gsp", self.ps) # TODO - if self.pi is not None: - write_height_field("pri_gsp", self.pi) # TODO - if self.pg is not None: - write_height_field("prg_gsp", self.pg) # TODO - if self.pre is not None: - write_height_field("pre_gsp", self.pre) # TODO - - -def setup_graupel(inp: GraupelInput, dt: float, qnc: float, backend: model_backends.BackendLike): - with utils.recursion_limit(10**4): # TODO thread safe? +def setup_graupel( + inp: common.GraupelInput, dt: float, qnc: float, backend: model_backends.BackendLike +): + with utils.recursion_limit(10**4): # TODO(havogt): make an option in gt4py? graupel_run_program = model_options.setup_program( backend=backend, program=graupel.graupel_run, @@ -230,7 +55,6 @@ def setup_graupel(inp: GraupelInput, dt: float, qnc: float, backend: model_backe vertical_sizes={ "vertical_start": gtx.int32(0), "vertical_end": gtx.int32(inp.nlev), - "last_lev": gtx.int32(inp.nlev - 1), }, offset_provider={"Koff": dims.KDim}, ) @@ -244,8 +68,8 @@ def main(): backend = model_backends.BACKENDS[args.backend] allocator = model_backends.get_allocator(backend) - inp = GraupelInput.load(filename=pathlib.Path(args.input_file), allocator=allocator) - out = GraupelOutput.allocate( + inp = common.GraupelInput.load(filename=pathlib.Path(args.input_file), allocator=allocator) + out = common.GraupelOutput.allocate( domain=gtx.domain({dims.CellDim: inp.ncells, dims.KDim: inp.nlev}), allocator=allocator ) @@ -262,19 +86,9 @@ def main(): te=inp.t, p=inp.p, rho=inp.rho, - qve=inp.qv, - qce=inp.qc, - qre=inp.qr, - qse=inp.qs, - qie=inp.qi, - qge=inp.qg, + q_in=inp.q, t_out=out.t, - qv_out=out.qv, - qc_out=out.qc, - qr_out=out.qr, - qs_out=out.qs, - qi_out=out.qi, - qg_out=out.qg, + q_out=out.q, pflx=out.pflx, pr=out.pr, ps=out.ps, diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py index 31717fa2b8..2bdcd287c6 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py @@ -5,11 +5,14 @@ # # Please, refer to the LICENSE file in the root directory. # SPDX-License-Identifier: BSD-3-Clause +from typing import NamedTuple + import gt4py.next as gtx from gt4py.next import maximum, minimum, power, sqrt, where from gt4py.next.experimental import concat_where from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.common.frozen import g_ct, idx, t_d +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import Q from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.properties import ( _deposition_auto_conversion, _deposition_factor, @@ -49,42 +52,74 @@ from icon4py.model.common.dimension import Koff -# @gtx.scan_operator(axis=dims.KDim, forward=True, init=(0.0, 0.0, 0.0, False)) <=== this should work now -@gtx.scan_operator(axis=dims.KDim, forward=True, init=(0.0, 0.0, 0.0, False)) +class PrecipState(NamedTuple): + q_update: ta.wpfloat + flx: ta.wpfloat + rho: ta.wpfloat + vc: ta.wpfloat + activated: bool + + +@gtx.scan_operator( + axis=dims.KDim, + forward=True, + init=PrecipState( + q_update=0.0, + flx=0.0, + rho=0.0, + vc=0.0, + activated=False, + ), +) def _precip( - state: tuple[ta.wpfloat, ta.wpfloat, ta.wpfloat, bool], + previous_level: PrecipState, prefactor: ta.wpfloat, # param[0] of fall_speed exponent: ta.wpfloat, # param[1] of fall_speed offset: ta.wpfloat, # param[1] of fall_speed zeta: ta.wpfloat, # dt/(2dz) vc: ta.wpfloat, # state dependent fall speed correction q: ta.wpfloat, # specific mass of hydrometeor - q_kp1: ta.wpfloat, # specific mass in next lower cell rho: ta.wpfloat, # density - mask: bool, # k-level located in cloud -) -> tuple[ta.wpfloat, ta.wpfloat, ta.wpfloat, bool]: # updates - _, flx, vt, is_level_activated = state - is_level_activated = is_level_activated | mask + mask: bool, +) -> PrecipState: + current_level_activated = previous_level.activated | mask rho_x = q * rho - flx_eff = (rho_x / zeta) + 2.0 * flx + flx_eff = (rho_x / zeta) + 2.0 * previous_level.flx # Inlined calculation using _fall_speed_scalar flx_partial = minimum(rho_x * vc * prefactor * power((rho_x + offset), exponent), flx_eff) - if is_level_activated: - update0 = (zeta * (flx_eff - flx_partial)) / ((1.0 + zeta * vt) * rho) # q update - update1 = (update0 * rho * vt + flx_partial) * 0.5 # flux - rho_x = (update0 + q_kp1) * 0.5 * rho - # Inlined calculation using _fall_speed_scalar - update2 = vc * prefactor * power((rho_x + offset), exponent) # vt + + rhox_prev = (previous_level.q_update + q) * 0.5 * previous_level.rho + + if previous_level.activated: + # this looks weird because we are setting vt based on previous level being active. bug? + vt = previous_level.vc * prefactor * power((rhox_prev + offset), exponent) + else: + vt = 0.0 + + if current_level_activated: + next_q_update = (zeta * (flx_eff - flx_partial)) / ((1.0 + zeta * vt) * rho) # q update + next_flx = (next_q_update * rho * vt + flx_partial) * 0.5 # flux else: - update0 = q - update1 = 0.0 - update2 = 0.0 - return update0, update1, update2, is_level_activated + next_q_update = q + next_flx = 0.0 + return PrecipState( + q_update=next_q_update, + flx=next_flx, + rho=rho, + vc=vc, + activated=current_level_activated, + ) + +class TempState(NamedTuple): + t: ta.wpfloat + eflx: ta.wpfloat + activated: bool -@gtx.scan_operator(axis=dims.KDim, forward=True, init=(0.0, 0.0, False)) + +@gtx.scan_operator(axis=dims.KDim, forward=True, init=TempState(t=0.0, eflx=0.0, activated=False)) def _temperature_update( - state: tuple[ta.wpfloat, ta.wpfloat, bool], + previous_level: TempState, t: ta.wpfloat, t_kp1: ta.wpfloat, ei_old: ta.wpfloat, @@ -97,17 +132,14 @@ def _temperature_update( dz: ta.wpfloat, dt: ta.wpfloat, mask: bool, -) -> tuple[ta.wpfloat, ta.wpfloat, bool]: - _, eflx, is_level_activated = state - is_level_activated = is_level_activated | mask - if is_level_activated: - e_int = ei_old + eflx - +) -> TempState: + current_level_activated = previous_level.activated | mask + if current_level_activated: eflx = dt * ( pr * (t_d.clw * t - t_d.cvd * t_kp1 - g_ct.lvc) + (pflx_tot) * (g_ct.ci * t - t_d.cvd * t_kp1 - g_ct.lsc) ) - e_int = e_int - eflx + e_int = ei_old + previous_level.eflx - eflx # Inlined calculation using T_from_internal_energy_scalar # in order to avoid scan_operator -> field_operator @@ -116,40 +148,10 @@ def _temperature_update( (t_d.cvd * (1.0 - qtot) + t_d.cvv * qv + t_d.clw * qliq + g_ct.ci * qice) * rho * dz ) # Moist isometric specific heat t = (e_int + rho * dz * (qliq * g_ct.lvc + qice * g_ct.lsc)) / cv + else: + eflx = previous_level.eflx - return t, eflx, is_level_activated - - -@gtx.field_operator -def _graupel_mask( - t: fa.CellKField[ta.wpfloat], # Temperature - rho: fa.CellKField[ta.wpfloat], # Density - qv: fa.CellKField[ta.wpfloat], # Q vapor content - qc: fa.CellKField[ta.wpfloat], # Q cloud content - qg: fa.CellKField[ta.wpfloat], # Q graupel content - qi: fa.CellKField[ta.wpfloat], # Q ice content - qr: fa.CellKField[ta.wpfloat], # Q rain content - qs: fa.CellKField[ta.wpfloat], # Q snow content -) -> tuple[ - fa.CellKField[bool], - fa.CellKField[bool], - fa.CellKField[bool], - fa.CellKField[bool], - fa.CellKField[bool], - fa.CellKField[bool], -]: - mask = where( - (maximum(qc, maximum(qg, maximum(qi, maximum(qr, qs)))) > g_ct.qmin) - | ((t < g_ct.tfrz_het2) & (qv > _qsat_ice_rho(t, rho))), - True, - False, - ) - is_sig_present = maximum(qg, maximum(qi, qs)) > g_ct.qmin - kmin_r = where(qr > g_ct.qmin, True, False) - kmin_i = where(qi > g_ct.qmin, True, False) - kmin_s = where(qs > g_ct.qmin, True, False) - kmin_g = where(qg > g_ct.qmin, True, False) - return mask, is_sig_present, kmin_r, kmin_i, kmin_s, kmin_g + return TempState(t=t, eflx=eflx, activated=current_level_activated) @gtx.field_operator @@ -157,53 +159,48 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] t: fa.CellKField[ta.wpfloat], p: fa.CellKField[ta.wpfloat], rho: fa.CellKField[ta.wpfloat], - qv: fa.CellKField[ta.wpfloat], # Q vapor content - qc: fa.CellKField[ta.wpfloat], # Q cloud content - qr: fa.CellKField[ta.wpfloat], # Q rain content - qs: fa.CellKField[ta.wpfloat], # Q snow content - qi: fa.CellKField[ta.wpfloat], # Q ice content - qg: fa.CellKField[ta.wpfloat], # Q graupel content - mask: fa.CellKField[bool], - is_sig_present: fa.CellKField[bool], + q: Q, dt: ta.wpfloat, qnc: ta.wpfloat, ) -> tuple[ - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], + Q, fa.CellKField[ta.wpfloat], ]: - dvsw = qv - _qsat_rho(t, rho) - qvsi = _qsat_ice_rho(t, rho) - dvsi = qv - qvsi + mask = where( + (maximum(q.c, maximum(q.g, maximum(q.i, maximum(q.r, q.s)))) > g_ct.qmin) + | ((t < g_ct.tfrz_het2) & (q.v > _qsat_ice_rho(t, rho))), + True, + False, + ) + is_sig_present = maximum(q.g, maximum(q.i, q.s)) > g_ct.qmin - n_snow = _snow_number(t, rho, qs) + dvsw = q.v - _qsat_rho(t, rho) + qvsi = _qsat_ice_rho(t, rho) + dvsi = q.v - qvsi + n_snow = _snow_number(t, rho, q.s) - l_snow = _snow_lambda(rho, qs, n_snow) + l_snow = _snow_lambda(rho, q.s, n_snow) # Define conversion 'matrix' - sx2x_c_r = _cloud_to_rain(t, qc, qr, qnc) - sx2x_r_v = _rain_to_vapor(t, rho, qc, qr, dvsw, dt) - sx2x_c_i = _cloud_x_ice(t, qc, qi, dt) + sx2x_c_r = _cloud_to_rain(t, q.c, q.r, qnc) + sx2x_r_v = _rain_to_vapor(t, rho, q.c, q.r, dvsw, dt) + sx2x_c_i = _cloud_x_ice(t, q.c, q.i, dt) sx2x_i_c = -minimum(sx2x_c_i, 0.0) sx2x_c_i = maximum(sx2x_c_i, 0.0) - sx2x_c_s = _cloud_to_snow(t, qc, qs, n_snow, l_snow) - sx2x_c_g = _cloud_to_graupel(t, rho, qc, qg) + sx2x_c_s = _cloud_to_snow(t, q.c, q.s, n_snow, l_snow) + sx2x_c_g = _cloud_to_graupel(t, rho, q.c, q.g) t_below_tmelt = t < t_d.tmelt - t_at_least_tmelt = not t_below_tmelt + t_at_least_tmelt = ~t_below_tmelt - n_ice = where(t_below_tmelt, _ice_number(t, rho), 0.0) - m_ice = where(t_below_tmelt, _ice_mass(qi, n_ice), 0.0) - x_ice = where(t_below_tmelt, _ice_sticking(t), 0.0) + n_ice = _ice_number(t, rho) + m_ice = _ice_mass(q.i, n_ice) + x_ice = _ice_sticking(t) eta = where(t_below_tmelt & is_sig_present, _deposition_factor(t, qvsi), 0.0) sx2x_v_i = where( - t_below_tmelt & is_sig_present, _vapor_x_ice(qi, m_ice, eta, dvsi, rho, dt), 0.0 + t_below_tmelt & is_sig_present, _vapor_x_ice(q.i, m_ice, eta, dvsi, rho, dt), 0.0 ) sx2x_i_v = where(t_below_tmelt & is_sig_present, -minimum(sx2x_v_i, 0.0), 0.0) sx2x_v_i = where(t_below_tmelt & is_sig_present, maximum(sx2x_v_i, 0.0), sx2x_i_v) @@ -212,19 +209,21 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] # TODO(): _deposition_auto_conversion yields roundoff differences in sx2x_i_s sx2x_i_s = where( t_below_tmelt & is_sig_present, - _deposition_auto_conversion(qi, m_ice, ice_dep) + _ice_to_snow(qi, n_snow, l_snow, x_ice), + _deposition_auto_conversion(q.i, m_ice, ice_dep) + _ice_to_snow(q.i, n_snow, l_snow, x_ice), 0.0, ) - sx2x_i_g = where(t_below_tmelt & is_sig_present, _ice_to_graupel(rho, qr, qg, qi, x_ice), 0.0) - sx2x_s_g = where(t_below_tmelt & is_sig_present, _snow_to_graupel(t, rho, qc, qs), 0.0) + sx2x_i_g = where( + t_below_tmelt & is_sig_present, _ice_to_graupel(rho, q.r, q.g, q.i, x_ice), 0.0 + ) + sx2x_s_g = where(t_below_tmelt & is_sig_present, _snow_to_graupel(t, rho, q.c, q.s), 0.0) sx2x_r_g = where( t_below_tmelt & is_sig_present, - _rain_to_graupel(t, rho, qc, qr, qi, qs, m_ice, dvsw, dt), + _rain_to_graupel(t, rho, q.c, q.r, q.i, q.s, m_ice, dvsw, dt), 0.0, ) sx2x_v_i = where( - t_below_tmelt, sx2x_v_i + _ice_deposition_nucleation(t, qc, qi, n_ice, dvsi, dt), 0.0 + t_below_tmelt, sx2x_v_i + _ice_deposition_nucleation(t, q.c, q.i, n_ice, dvsi, dt), 0.0 ) # 0.0 or sx2x_v_i both OK sx2x_c_r = where(t_at_least_tmelt, sx2x_c_r + sx2x_c_s + sx2x_c_g, sx2x_c_r) sx2x_c_s = where(t_at_least_tmelt, 0.0, sx2x_c_s) @@ -232,21 +231,21 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] ice_dep = where(t_at_least_tmelt, 0.0, ice_dep) eta = where(t_at_least_tmelt, 0.0, eta) - dvsw0 = where(is_sig_present, qv - _qsat_rho_tmelt(rho), 0.0) + dvsw0 = where(is_sig_present, q.v - _qsat_rho_tmelt(rho), 0.0) sx2x_v_s = where( is_sig_present, - _vapor_x_snow(t, p, rho, qs, n_snow, l_snow, eta, ice_dep, dvsw, dvsi, dvsw0, dt), + _vapor_x_snow(t, p, rho, q.s, n_snow, l_snow, eta, ice_dep, dvsw, dvsi, dvsw0, dt), 0.0, ) sx2x_s_v = where(is_sig_present, -minimum(sx2x_v_s, 0.0), 0.0) sx2x_v_s = where(is_sig_present, maximum(sx2x_v_s, 0.0), 0.0) - sx2x_v_g = where(is_sig_present, _vapor_x_graupel(t, p, rho, qg, dvsw, dvsi, dvsw0, dt), 0.0) + sx2x_v_g = where(is_sig_present, _vapor_x_graupel(t, p, rho, q.g, dvsw, dvsi, dvsw0, dt), 0.0) sx2x_g_v = where(is_sig_present, -minimum(sx2x_v_g, 0.0), 0.0) sx2x_v_g = where(is_sig_present, maximum(sx2x_v_g, 0.0), 0.0) - sx2x_s_r = where(is_sig_present, _snow_to_rain(t, p, rho, dvsw0, qs), 0.0) - sx2x_g_r = where(is_sig_present, _graupel_to_rain(t, p, rho, dvsw0, qg), 0.0) + sx2x_s_r = where(is_sig_present, _snow_to_rain(t, p, rho, dvsw0, q.s), 0.0) + sx2x_g_r = where(is_sig_present, _graupel_to_rain(t, p, rho, dvsw0, q.g), 0.0) # The following transitions are not physically meaningful, would be 0.0 in other implementation # here they are simply never used: @@ -270,8 +269,8 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] ) # Missing: sx2x_g_c + sx2x_g_s + sx2x_g_i # if ((sink[qx_ind[ix]] > stot) && (q[qx_ind[ix]].x[oned_vec_index] > qmin)) - stot = qv / dt - sink_v_saturated = (sink_v > stot) & (qv > g_ct.qmin) + stot = q.v / dt + sink_v_saturated = (sink_v > stot) & (q.v > g_ct.qmin) sx2x_v_s = where(sink_v_saturated, sx2x_v_s * stot / sink_v, sx2x_v_s) sx2x_v_i = where(sink_v_saturated, sx2x_v_i * stot / sink_v, sx2x_v_i) sx2x_v_g = where(sink_v_saturated, sx2x_v_g * stot / sink_v, sx2x_v_g) @@ -279,8 +278,8 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] sink_v_saturated, sx2x_v_s + sx2x_v_i + sx2x_v_g, sink_v ) # Missing: sx2x_v_c + sx2x_v_r - stot = qc / dt - sink_c_saturated = (sink_c > stot) & (qc > g_ct.qmin) + stot = q.c / dt + sink_c_saturated = (sink_c > stot) & (q.c > g_ct.qmin) sx2x_c_r = where(sink_c_saturated, sx2x_c_r * stot / sink_c, sx2x_c_r) sx2x_c_s = where(sink_c_saturated, sx2x_c_s * stot / sink_c, sx2x_c_s) sx2x_c_i = where(sink_c_saturated, sx2x_c_i * stot / sink_c, sx2x_c_i) @@ -289,16 +288,16 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] sink_c_saturated, sx2x_c_r + sx2x_c_s + sx2x_c_i + sx2x_c_g, sink_c ) # Missing: sx2x_c_v - stot = qr / dt - sink_r_saturated = (sink_r > stot) & (qr > g_ct.qmin) + stot = q.r / dt + sink_r_saturated = (sink_r > stot) & (q.r > g_ct.qmin) sx2x_r_v = where(sink_r_saturated, sx2x_r_v * stot / sink_r, sx2x_r_v) sx2x_r_g = where(sink_r_saturated, sx2x_r_g * stot / sink_r, sx2x_r_g) sink_r = where( sink_r_saturated, sx2x_r_v + sx2x_r_g, sink_r ) # Missing: sx2x_r_c + sx2x_r_s + sx2x_r_i - stot = qs / dt - sink_s_saturated = (sink_s > stot) & (qs > g_ct.qmin) + stot = q.s / dt + sink_s_saturated = (sink_s > stot) & (q.s > g_ct.qmin) sx2x_s_v = where(sink_s_saturated, sx2x_s_v * stot / sink_s, sx2x_s_v) sx2x_s_r = where(sink_s_saturated, sx2x_s_r * stot / sink_s, sx2x_s_r) sx2x_s_g = where(sink_s_saturated, sx2x_s_g * stot / sink_s, sx2x_s_g) @@ -306,8 +305,8 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] sink_s_saturated, sx2x_s_v + sx2x_s_r + sx2x_s_g, sink_s ) # Missing: sx2x_s_c + sx2x_s_i - stot = qi / dt - sink_i_saturated = (sink_i > stot) & (qi > g_ct.qmin) + stot = q.i / dt + sink_i_saturated = (sink_i > stot) & (q.i > g_ct.qmin) sx2x_i_v = where(sink_i_saturated, sx2x_i_v * stot / sink_i, sx2x_i_v) sx2x_i_c = where(sink_i_saturated, sx2x_i_c * stot / sink_i, sx2x_i_c) sx2x_i_s = where(sink_i_saturated, sx2x_i_s * stot / sink_i, sx2x_i_s) @@ -316,8 +315,8 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] sink_i_saturated, sx2x_i_v + sx2x_i_c + sx2x_i_s + sx2x_i_g, sink_i ) # Missing: sx2x_i_r - stot = qg / dt - sink_g_saturated = (sink_g > stot) & (qg > g_ct.qmin) + stot = q.g / dt + sink_g_saturated = (sink_g > stot) & (q.g > g_ct.qmin) sx2x_g_v = where(sink_g_saturated, sx2x_g_v * stot / sink_g, sx2x_g_v) sx2x_g_r = where(sink_g_saturated, sx2x_g_r * stot / sink_g, sx2x_g_r) sink_g = where( @@ -327,17 +326,17 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] # water content updates: # Physical: v_s, v_i, v_g, c_r, c_s, c_i, c_g, r_v, r_g, s_v, s_r, s_g, i_v, i_c, i_s, i_g, g_v, g_r dqdt_v = sx2x_r_v + sx2x_s_v + sx2x_i_v + sx2x_g_v - sink_v # Missing: sx2x_c_v - qv = where(mask, maximum(0.0, qv + dqdt_v * dt), qv) + qv = where(mask, maximum(0.0, q.v + dqdt_v * dt), q.v) dqdt_c = sx2x_i_c - sink_c # Missing: sx2x_v_c, sx2x_r_c, sx2x_s_c, sx2x_g_c - qc = where(mask, maximum(0.0, qc + dqdt_c * dt), qc) + qc = where(mask, maximum(0.0, q.c + dqdt_c * dt), q.c) dqdt_r = sx2x_c_r + sx2x_s_r + sx2x_g_r - sink_r # Missing: sx2x_v_r + sx2x_i_r - qr = where(mask, maximum(0.0, qr + dqdt_r * dt), qr) + qr = where(mask, maximum(0.0, q.r + dqdt_r * dt), q.r) dqdt_s = sx2x_v_s + sx2x_c_s + sx2x_i_s - sink_s # Missing: sx2x_r_s + sx2x_g_s - qs = where(mask, maximum(0.0, qs + dqdt_s * dt), qs) + qs = where(mask, maximum(0.0, q.s + dqdt_s * dt), q.s) dqdt_i = sx2x_v_i + sx2x_c_i - sink_i # Missing: sx2x_r_i + sx2x_s_i + sx2x_g_i - qi = where(mask, maximum(0.0, qi + dqdt_i * dt), qi) + qi = where(mask, maximum(0.0, q.i + dqdt_i * dt), q.i) dqdt_g = sx2x_v_g + sx2x_c_g + sx2x_r_g + sx2x_s_g + sx2x_i_g - sink_g - qg = where(mask, maximum(0.0, qg + dqdt_g * dt), qg) + qg = where(mask, maximum(0.0, q.g + dqdt_g * dt), q.g) qice = qs + qi + qg qliq = qc + qr @@ -360,7 +359,7 @@ def _q_t_update( # noqa: PLR0915 [too-many-statements] / cv, t, ) - return qv, qc, qr, qs, qi, qg, t + return Q(v=qv, c=qc, r=qr, s=qs, i=qi, g=qg), t @gtx.field_operator @@ -370,12 +369,7 @@ def _precipitation_effects( kmin_i: fa.CellKField[bool], # ice minimum level kmin_s: fa.CellKField[bool], # snow minimum level kmin_g: fa.CellKField[bool], # graupel minimum level - qv: fa.CellKField[ta.wpfloat], # Q vapor content - qc: fa.CellKField[ta.wpfloat], # Q cloud content - qr: fa.CellKField[ta.wpfloat], # Q rain content - qs: fa.CellKField[ta.wpfloat], # Q snow content - qi: fa.CellKField[ta.wpfloat], # Q ice content - qg: fa.CellKField[ta.wpfloat], # Q graupel content qv, + q_in: Q, t: fa.CellKField[ta.wpfloat], # temperature, rho: fa.CellKField[ta.wpfloat], # density dz: fa.CellKField[ta.wpfloat], @@ -394,69 +388,55 @@ def _precipitation_effects( fa.CellKField[ta.wpfloat], ]: # Store current fields for later temperature update - qliq = qc + qr - qice = qs + qi + qg - ei_old = _internal_energy(t, qv, qliq, qice, rho, dz) + qliq = q_in.c + q_in.r + qice = q_in.s + q_in.i + q_in.g + ei_old = _internal_energy(t, q_in.v, qliq, qice, rho, dz) zeta = dt / (2.0 * dz) xrho = sqrt(g_ct.rho_00 / rho) vc_r = _vel_scale_factor_default(xrho) - vc_s = _vel_scale_factor_snow(xrho, rho, t, qs) + vc_s = _vel_scale_factor_snow(xrho, rho, t, q_in.s) vc_i = _vel_scale_factor_ice(xrho) vc_g = _vel_scale_factor_default(xrho) - q_kp1 = concat_where(dims.KDim < last_lev, qr(Koff[1]), qr) - qr, pr, _, _ = _precip( - idx.prefactor_r, idx.exponent_r, idx.offset_r, zeta, vc_r, qr, q_kp1, rho, kmin_r + qr, pr, _, _, _ = _precip( + idx.prefactor_r, idx.exponent_r, idx.offset_r, zeta, vc_r, q_in.r, rho, kmin_r ) - q_kp1 = concat_where(dims.KDim < last_lev, qs(Koff[1]), qs) - qs, ps, _, _ = _precip( - idx.prefactor_s, idx.exponent_s, idx.offset_s, zeta, vc_s, qs, q_kp1, rho, kmin_s + qs, ps, _, _, _ = _precip( + idx.prefactor_s, idx.exponent_s, idx.offset_s, zeta, vc_s, q_in.s, rho, kmin_s ) - q_kp1 = concat_where(dims.KDim < last_lev, qi(Koff[1]), qi) - qi, pi, _, _ = _precip( - idx.prefactor_i, idx.exponent_i, idx.offset_i, zeta, vc_i, qi, q_kp1, rho, kmin_i + qi, pi, _, _, _ = _precip( + idx.prefactor_i, idx.exponent_i, idx.offset_i, zeta, vc_i, q_in.i, rho, kmin_i ) - q_kp1 = concat_where(dims.KDim < last_lev, qg(Koff[1]), qg) - qg, pg, _, _ = _precip( - idx.prefactor_g, idx.exponent_g, idx.offset_g, zeta, vc_g, qg, q_kp1, rho, kmin_g + qg, pg, _, _, _ = _precip( + idx.prefactor_g, idx.exponent_g, idx.offset_g, zeta, vc_g, q_in.g, rho, kmin_g ) - qliq = qc + qr + qliq = q_in.c + qr qice = qs + qi + qg - p_sig = ps + pi + pg + pflx_tot = ps + pi + pg t_kp1 = concat_where(dims.KDim < last_lev, t(Koff[1]), t) kmin_rsig = kmin_r | kmin_s | kmin_i | kmin_g t, eflx, _ = _temperature_update( - t, t_kp1, ei_old, pr, p_sig, qv, qliq, qice, rho, dz, dt, kmin_rsig + t, t_kp1, ei_old, pr, pflx_tot, q_in.v, qliq, qice, rho, dz, dt, kmin_rsig ) - return qr, qs, qi, qg, t, p_sig + pr, pr, ps, pi, pg, eflx / dt + return qr, qs, qi, qg, t, pflx_tot + pr, pr, ps, pi, pg, eflx / dt @gtx.field_operator -def _graupel_run( - last_lev: gtx.int32, +def graupel( + last_level: gtx.int32, dz: fa.CellKField[ta.wpfloat], te: fa.CellKField[ta.wpfloat], # Temperature p: fa.CellKField[ta.wpfloat], # Pressure rho: fa.CellKField[ta.wpfloat], # Density containing dry air and water constituents - qve: fa.CellKField[ta.wpfloat], # Specific humidity - qce: fa.CellKField[ta.wpfloat], # Specific cloud water content - qre: fa.CellKField[ta.wpfloat], # Specific rain water - qse: fa.CellKField[ta.wpfloat], # Specific snow water - qie: fa.CellKField[ta.wpfloat], # Specific ice water content - qge: fa.CellKField[ta.wpfloat], # Specific graupel water content + q: Q, dt: ta.wpfloat, qnc: ta.wpfloat, ) -> tuple[ fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], - fa.CellKField[ta.wpfloat], + Q, fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], @@ -464,41 +444,29 @@ def _graupel_run( fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat], ]: - mask, is_sig_present, kmin_r, kmin_i, kmin_s, kmin_g = _graupel_mask( - te, rho, qve, qce, qge, qie, qre, qse - ) - qv, qc, qr, qs, qi, qg, t = _q_t_update( - te, p, rho, qve, qce, qre, qse, qie, qge, mask, is_sig_present, dt, qnc - ) + kmin_r = where(q.r > g_ct.qmin, True, False) + kmin_i = where(q.i > g_ct.qmin, True, False) + kmin_s = where(q.s > g_ct.qmin, True, False) + kmin_g = where(q.g > g_ct.qmin, True, False) + q, t = _q_t_update(te, p, rho, q, dt, qnc) qr, qs, qi, qg, t, pflx, pr, ps, pi, pg, pre = _precipitation_effects( - last_lev, kmin_r, kmin_i, kmin_s, kmin_g, qv, qc, qr, qs, qi, qg, t, rho, dz, dt + last_level, kmin_r, kmin_i, kmin_s, kmin_g, q, t, rho, dz, dt ) - return t, qv, qc, qr, qs, qi, qg, pflx, pr, ps, pi, pg, pre + return t, Q(v=q.v, c=q.c, r=qr, s=qs, i=qi, g=qg), pflx, pr, ps, pi, pg, pre @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) def graupel_run( - last_lev: gtx.int32, dz: fa.CellKField[ta.wpfloat], te: fa.CellKField[ta.wpfloat], # Temperature p: fa.CellKField[ta.wpfloat], # Pressure rho: fa.CellKField[ta.wpfloat], # Density containing dry air and water constituents - qve: fa.CellKField[ta.wpfloat], # Specific humidityn - qce: fa.CellKField[ta.wpfloat], # Specific cloud water content - qre: fa.CellKField[ta.wpfloat], # Specific rain water - qse: fa.CellKField[ta.wpfloat], # Specific snow water - qie: fa.CellKField[ta.wpfloat], # Specific ice water content - qge: fa.CellKField[ta.wpfloat], # Specific graupel water content + q_in: Q, dt: ta.wpfloat, # Time step qnc: ta.wpfloat, + q_out: Q, t_out: fa.CellKField[ta.wpfloat], # Revised temperature - qv_out: fa.CellKField[ta.wpfloat], # Revised humidity - qc_out: fa.CellKField[ta.wpfloat], # Revised cloud water - qr_out: fa.CellKField[ta.wpfloat], # Revised rain water - qs_out: fa.CellKField[ta.wpfloat], # Revised snow water - qi_out: fa.CellKField[ta.wpfloat], # Revised ice water - qg_out: fa.CellKField[ta.wpfloat], # Revised graupel water pflx: fa.CellKField[ta.wpfloat], # Total precipitation flux pr: fa.CellKField[ta.wpfloat], # Precipitation of rain ps: fa.CellKField[ta.wpfloat], # Precipitation of snow @@ -510,21 +478,16 @@ def graupel_run( vertical_start: gtx.int32, vertical_end: gtx.int32, ): - _graupel_run( - last_lev, - dz, - te, - p, - rho, - qve, - qce, - qre, - qse, - qie, - qge, - dt, - qnc, - out=(t_out, qv_out, qc_out, qr_out, qs_out, qi_out, qg_out, pflx, pr, ps, pi, pg, pre), + graupel( + last_level=vertical_end - 1, + dz=dz, + te=te, + p=p, + rho=rho, + q=q_in, + dt=dt, + qnc=qnc, + out=(t_out, q_out, pflx, pr, ps, pi, pg, pre), domain={ dims.CellDim: (horizontal_start, horizontal_end), dims.KDim: (vertical_start, vertical_end), diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/muphys.py new file mode 100644 index 0000000000..1e7f64e48c --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/muphys.py @@ -0,0 +1,101 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import gt4py.next as gtx + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import Q +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.saturation_adjustment import ( + _saturation_adjustment, +) +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations.graupel import graupel +from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta + + +@gtx.field_operator +def _muphys( + last_level: gtx.int32, + dz: fa.CellKField[ta.wpfloat], + te: fa.CellKField[ta.wpfloat], # Temperature + p: fa.CellKField[ta.wpfloat], # Pressure + rho: fa.CellKField[ta.wpfloat], # Density containing dry air and water constituents + q_in: Q, + dt: ta.wpfloat, + qnc: ta.wpfloat, +) -> tuple[ + fa.CellKField[ta.wpfloat], + Q, + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], + fa.CellKField[ta.wpfloat], +]: + te, qve, qce = _saturation_adjustment( + te=te, + q_in=q_in, + rho=rho, + ) + + t, q, pflx, pr, ps, pi, pg, pre = graupel( + last_level, + dz, + te, + p, + rho, + Q(v=qve, c=qce, r=q_in.r, s=q_in.s, i=q_in.i, g=q_in.g), + dt, + qnc, + ) + + te, qve, qce = _saturation_adjustment( + te=t, + q_in=q, + rho=rho, + ) + + return t, Q(v=qve, c=qce, r=q.r, s=q.s, i=q.i, g=q.g), pflx, pr, ps, pi, pg, pre + + +@gtx.program(grid_type=gtx.GridType.UNSTRUCTURED) +def muphys_run( + dz: fa.CellKField[ta.wpfloat], + te: fa.CellKField[ta.wpfloat], # Temperature + p: fa.CellKField[ta.wpfloat], # Pressure + rho: fa.CellKField[ta.wpfloat], # Density containing dry air and water constituents + q_in: Q, + dt: ta.wpfloat, # Time step + qnc: ta.wpfloat, + q_out: Q, + t_out: fa.CellKField[ta.wpfloat], # Revised temperature + pflx: fa.CellKField[ta.wpfloat], # Total precipitation flux + pr: fa.CellKField[ta.wpfloat], # Precipitation of rain + ps: fa.CellKField[ta.wpfloat], # Precipitation of snow + pi: fa.CellKField[ta.wpfloat], # Precipitation of ice + pg: fa.CellKField[ta.wpfloat], # Precipitation of graupel + pre: fa.CellKField[ta.wpfloat], # Precipitation of graupel + horizontal_start: gtx.int32, + horizontal_end: gtx.int32, + vertical_start: gtx.int32, + vertical_end: gtx.int32, +): + _muphys( + vertical_end - 1, + dz, + te, + p, + rho, + q_in, + dt, + qnc, + out=(t_out, q_out, pflx, pr, ps, pi, pg, pre), + domain={ + dims.CellDim: (horizontal_start, horizontal_end), + dims.KDim: (vertical_start, vertical_end), + }, + ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_full_muphys.py new file mode 100644 index 0000000000..599016b038 --- /dev/null +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_full_muphys.py @@ -0,0 +1,139 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +from __future__ import annotations + +import dataclasses +import pathlib +from typing import Final + +import numpy as np +import pytest +from gt4py import next as gtx + +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import common, run_full_muphys +from icon4py.model.common import dimension as dims, model_backends +from icon4py.model.testing import data_handling, definitions as testing_defs +from icon4py.model.testing.fixtures.datatest import backend_like + + +def _path_to_experiment_testdata(experiment: MuphysGraupelExperiment) -> pathlib.Path: + return testing_defs.get_test_data_root_path() / "full_muphys_data" / experiment.name + + +@dataclasses.dataclass(frozen=True) +class MuphysGraupelExperiment: + name: str + uri: str + dtype: np.dtype + dt: float = 30.0 + qnc: float = 100.0 + + @property + def input_file(self) -> pathlib.Path: + return _path_to_experiment_testdata(self) / "input.nc" + + @property + def reference_file(self) -> pathlib.Path: + return _path_to_experiment_testdata(self) / "reference.nc" + + def __str__(self): + return self.name + + +class Experiments: + # TODO currently on havogt's polybox + MINI: Final = MuphysGraupelExperiment( + name="mini", + uri="TODO", + dtype=np.float32, + ) + TINY: Final = MuphysGraupelExperiment( + name="tiny", + uri="TODO", + dtype=np.float64, + ) + R2B05: Final = MuphysGraupelExperiment( + name="R2B05", + uri="TODO", + dtype=np.float32, + ) + + +@pytest.fixture(autouse=True) +def download_test_data(experiment: MuphysGraupelExperiment) -> None: + """Downloads test data for an experiment (implicit fixture).""" + data_handling.download_test_data(_path_to_experiment_testdata(experiment), uri=experiment.uri) + + +@pytest.mark.datatest +@pytest.mark.parametrize( + "experiment", + [ + Experiments.MINI, + Experiments.TINY, + Experiments.R2B05, + ], + ids=lambda exp: exp.name, +) +@pytest.mark.parametrize("single_program", [True, False], ids=lambda sp: f"single_program={sp}") +def test_full_muphys( + backend_like: model_backends.BackendLike, + experiment: MuphysGraupelExperiment, + single_program: bool, +) -> None: + # TODO why does it verify on the graupel only reference? + inp = common.GraupelInput.load( + filename=experiment.input_file, allocator=model_backends.get_allocator(backend_like) + ) + + muphys_program = run_full_muphys.setup_muphys( + inp, + dt=experiment.dt, + qnc=experiment.qnc, + backend=backend_like, + single_program=single_program, + ) + + out = common.GraupelOutput.allocate( + allocator=model_backends.get_allocator(backend_like), + domain=gtx.domain({dims.CellDim: inp.ncells, dims.KDim: inp.nlev}), + ) + + muphys_program( + dz=inp.dz, + te=inp.t, + p=inp.p, + rho=inp.rho, + q_in=inp.q, + t_out=out.t, + q_out=out.q, + pflx=out.pflx, + pr=out.pr, + ps=out.ps, + pi=out.pi, + pg=out.pg, + pre=out.pre, + ) + + ref = common.GraupelOutput.load( + filename=experiment.reference_file, allocator=model_backends.get_allocator(backend_like) + ) + + # TODO check tolerances + rtol = 1e-14 if experiment.dtype == np.float64 else 1e-7 + atol = 1e-16 if experiment.dtype == np.float64 else 1e-8 + # TODO we run the float32 input experiments with float64 + + np.testing.assert_allclose(ref.qv.asnumpy(), out.qv.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qc.asnumpy(), out.qc.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qi.asnumpy(), out.qi.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qr.asnumpy(), out.qr.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qs.asnumpy(), out.qs.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.qg.asnumpy(), out.qg.asnumpy(), atol=atol, rtol=rtol) + np.testing.assert_allclose(ref.t.asnumpy(), out.t.asnumpy(), atol=atol, rtol=rtol) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py index 11c1be1ff2..87a377cba7 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/integration_tests/test_graupel_only.py @@ -16,7 +16,8 @@ import pytest from gt4py import next as gtx -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import run_graupel_only +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.driver import common, run_graupel_only +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.implementations import graupel from icon4py.model.common import dimension as dims, model_backends from icon4py.model.testing import data_handling, definitions as testing_defs from icon4py.model.testing.fixtures.datatest import backend_like @@ -82,10 +83,9 @@ def download_test_data(experiment: MuphysGraupelExperiment) -> None: ids=lambda exp: exp.name, ) def test_graupel_only( - backend_like: model_backends.BackendLike, - experiment: MuphysGraupelExperiment, + backend_like: model_backends.BackendLike, experiment: MuphysGraupelExperiment ) -> None: - inp = run_graupel_only.GraupelInput.load( + inp = common.GraupelInput.load( filename=experiment.input_file, allocator=model_backends.get_allocator(backend_like) ) @@ -93,7 +93,7 @@ def test_graupel_only( inp, dt=experiment.dt, qnc=experiment.qnc, backend=backend_like ) - out = run_graupel_only.GraupelOutput.allocate( + out = common.GraupelOutput.allocate( allocator=model_backends.get_allocator(backend_like), domain=gtx.domain({dims.CellDim: inp.ncells, dims.KDim: inp.nlev}), ) @@ -103,19 +103,9 @@ def test_graupel_only( te=inp.t, p=inp.p, rho=inp.rho, - qve=inp.qv, - qce=inp.qc, - qre=inp.qr, - qse=inp.qs, - qie=inp.qi, - qge=inp.qg, + q_in=inp.q, t_out=out.t, - qv_out=out.qv, - qc_out=out.qc, - qr_out=out.qr, - qs_out=out.qs, - qi_out=out.qi, - qg_out=out.qg, + q_out=out.q, pflx=out.pflx, pr=out.pr, ps=out.ps, @@ -124,9 +114,8 @@ def test_graupel_only( pre=out.pre, ) - ref = run_graupel_only.GraupelOutput.load( - filename=experiment.reference_file, - allocator=model_backends.get_allocator(backend_like), + ref = common.GraupelOutput.load( + filename=experiment.reference_file, allocator=model_backends.get_allocator(backend_like) ) # TODO check tolerances diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py index 58b83e8f52..69c06c9e0b 100644 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment.py @@ -12,7 +12,10 @@ import numpy as np import pytest -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.thermo import saturation_adjustment +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.definitions import Q +from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.saturation_adjustment import ( + saturation_adjustment, +) from icon4py.model.common import dimension as dims from icon4py.model.common.type_alias import wpfloat from icon4py.model.common.utils import data_allocation as data_alloc @@ -23,7 +26,6 @@ from icon4py.model.common.grid import base as base_grid -# @pytest.mark.embedded_only class TestSaturationAdjustment(StencilTest): PROGRAM = saturation_adjustment OUTPUTS = ("te_out", "qve_out", "qce_out") @@ -32,13 +34,6 @@ class TestSaturationAdjustment(StencilTest): def reference( grid: base_grid.Grid, te: np.ndarray, - qve: np.ndarray, - qce: np.ndarray, - qre: np.ndarray, - qse: np.ndarray, - qie: np.ndarray, - qge: np.ndarray, - rho: np.ndarray, **kwargs, ) -> dict: return dict( @@ -53,23 +48,25 @@ def input_data(self, grid: base_grid.Grid) -> dict: te=data_alloc.constant_field( grid, 273.90911754406039, dims.CellDim, dims.KDim, dtype=wpfloat ), - qve=data_alloc.constant_field( - grid, 4.4913424511676030e-003, dims.CellDim, dims.KDim, dtype=wpfloat - ), - qce=data_alloc.constant_field( - grid, 6.0066941654987605e-013, dims.CellDim, dims.KDim, dtype=wpfloat - ), - qre=data_alloc.constant_field( - grid, 2.5939378002267028e-004, dims.CellDim, dims.KDim, dtype=wpfloat - ), - qse=data_alloc.constant_field( - grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat - ), - qie=data_alloc.constant_field( - grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat - ), - qge=data_alloc.constant_field( - grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat + q_in=Q( + v=data_alloc.constant_field( + grid, 4.4913424511676030e-003, dims.CellDim, dims.KDim, dtype=wpfloat + ), + c=data_alloc.constant_field( + grid, 6.0066941654987605e-013, dims.CellDim, dims.KDim, dtype=wpfloat + ), + r=data_alloc.constant_field( + grid, 2.5939378002267028e-004, dims.CellDim, dims.KDim, dtype=wpfloat + ), + s=data_alloc.constant_field( + grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat + ), + i=data_alloc.constant_field( + grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat + ), + g=data_alloc.constant_field( + grid, 3.582312533881839e-06, dims.CellDim, dims.KDim, dtype=wpfloat + ), ), rho=data_alloc.constant_field( grid, 1.1371657035251757, dims.CellDim, dims.KDim, dtype=wpfloat @@ -77,4 +74,8 @@ def input_data(self, grid: base_grid.Grid) -> dict: te_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), qve_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), qce_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), + horizontal_start=0, + horizontal_end=grid.num_cells, + vertical_start=0, + vertical_end=grid.num_levels, ) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment2.py b/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment2.py deleted file mode 100644 index b8b685e11b..0000000000 --- a/model/atmosphere/subgrid_scale_physics/muphys/tests/muphys/stencil_tests/test_saturation_adjustment2.py +++ /dev/null @@ -1,81 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022-2024, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# Please, refer to the LICENSE file in the root directory. -# SPDX-License-Identifier: BSD-3-Clause -import numpy as np -import pytest - -from icon4py.model.atmosphere.subgrid_scale_physics.muphys.core.saturation_adjustment2 import ( - saturation_adjustment2, -) -from icon4py.model.common import dimension as dims -from icon4py.model.common.type_alias import wpfloat -from icon4py.model.common.utils import data_allocation as data_alloc -from icon4py.model.testing.stencil_tests import StencilTest - - -@pytest.mark.embedded_only -class TestSaturationAdjustment(StencilTest): - PROGRAM = saturation_adjustment2 - OUTPUTS = ("te_out", "qve_out", "qce_out") - - @staticmethod - def reference( - grid, - te: np.ndarray, - qve: np.ndarray, - qce: np.ndarray, - qre: np.ndarray, - qti: np.ndarray, - cvc: np.ndarray, - ue: np.ndarray, - Tx_hold: np.ndarray, - Tx: np.ndarray, - rho: np.ndarray, - **kwargs, - ) -> dict: - return dict( - te_out=np.full(te.shape, 273.91226488486984), - qve_out=np.full(te.shape, 4.4903852062454690e-003), - qce_out=np.full(te.shape, 9.5724552280369163e-007), - ) - - @pytest.fixture - def input_data(self, grid): - return dict( - te=data_alloc.constant_field( - grid, 273.90911754406039, dims.CellDim, dims.KDim, dtype=wpfloat - ), - qve=data_alloc.constant_field( - grid, 4.4913424511676030e-003, dims.CellDim, dims.KDim, dtype=wpfloat - ), - qce=data_alloc.constant_field( - grid, 6.0066941654987605e-013, dims.CellDim, dims.KDim, dtype=wpfloat - ), - qre=data_alloc.constant_field( - grid, 2.5939378002267028e-004, dims.CellDim, dims.KDim, dtype=wpfloat - ), - qti=data_alloc.constant_field( - grid, 1.0746937601645517e-005, dims.CellDim, dims.KDim, dtype=wpfloat - ), - rho=data_alloc.constant_field( - grid, 1.1371657035251757, dims.CellDim, dims.KDim, dtype=wpfloat - ), - cvc=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), # Temporary - ue=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), # Temporary - Tx_hold=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, dtype=wpfloat - ), # Temporary - Tx=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), # Temporary - qx_hold=data_alloc.zero_field( - grid, dims.CellDim, dims.KDim, dtype=wpfloat - ), # Temporary - qx=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), # Temporary - dqx=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), # Temporary - te_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qve_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - qce_out=data_alloc.zero_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat), - ) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index dfeb0ec6bf..7d81cd8c8e 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -32,7 +32,10 @@ def allocate_data( ) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: _allocate_field = constructors.as_field.partial(allocator=allocator) # type:ignore[attr-defined] # TODO(havogt): check why it doesn't understand the fluid_partial input_data = { - k: tuple(_allocate_field(domain=field.domain, data=field.ndarray) for field in v) + # TODO: make this work for dataclass named collections as well + k: v.__class__( + *tuple(_allocate_field(domain=field.domain, data=field.ndarray) for field in v) + ) if isinstance(v, tuple) else _allocate_field(domain=v.domain, data=v.ndarray) if not gtx.is_scalar_type(v) and k != "domain" From 213e801f1de3c3a58c24a7f0bbfd1ea0bdc1faf2 Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 5 Dec 2025 16:01:58 +0100 Subject: [PATCH 20/23] point to gt4py main --- pyproject.toml | 2 +- uv.lock | 86 ++++++++++++++++++++++++++++++-------------------- 2 files changed, 52 insertions(+), 36 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f073985645..e32dd27599 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -355,7 +355,7 @@ url = "https://test.pypi.org/simple/" [tool.uv.sources] dace = {git = "https://github.com/GridTools/dace", tag = "__gt4py-next-integration_2025_11_05"} ghex = {git = "https://github.com/msimberg/GHEX.git", branch = "async-mpi"} -gt4py = {git = "https://github.com/havogt/gt4py", branch = "cse_in_fieldop_fusion"} +gt4py = {git = "https://github.com/GridTools/gt4py", branch = "main"} # gt4py = {index = "test.pypi"} icon4py-atmosphere-advection = {workspace = true} icon4py-atmosphere-diffusion = {workspace = true} diff --git a/uv.lock b/uv.lock index 6ad5418b1f..ceb084e033 100644 --- a/uv.lock +++ b/uv.lock @@ -165,7 +165,7 @@ wheels = [ [[package]] name = "black" -version = "24.10.0" +version = "25.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -173,28 +173,33 @@ dependencies = [ { name = "packaging" }, { name = "pathspec" }, { name = "platformdirs" }, + { name = "pytokens" }, { name = "tomli", marker = "python_full_version < '3.11'" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d8/0d/cc2fb42b8c50d80143221515dd7e4766995bd07c56c9a3ed30baf080b6dc/black-24.10.0.tar.gz", hash = "sha256:846ea64c97afe3bc677b761787993be4991810ecc7a4a937816dd6bddedc4875", size = 645813, upload-time = "2024-10-07T19:20:50.361Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a3/f3/465c0eb5cddf7dbbfe1fecd9b875d1dcf51b88923cd2c1d7e9ab95c6336b/black-24.10.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e6668650ea4b685440857138e5fe40cde4d652633b1bdffc62933d0db4ed9812", size = 1623211, upload-time = "2024-10-07T19:26:12.43Z" }, - { url = "https://files.pythonhosted.org/packages/df/57/b6d2da7d200773fdfcc224ffb87052cf283cec4d7102fab450b4a05996d8/black-24.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1c536fcf674217e87b8cc3657b81809d3c085d7bf3ef262ead700da345bfa6ea", size = 1457139, upload-time = "2024-10-07T19:25:06.453Z" }, - { url = "https://files.pythonhosted.org/packages/6e/c5/9023b7673904a5188f9be81f5e129fff69f51f5515655fbd1d5a4e80a47b/black-24.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:649fff99a20bd06c6f727d2a27f401331dc0cc861fb69cde910fe95b01b5928f", size = 1753774, upload-time = "2024-10-07T19:23:58.47Z" }, - { url = "https://files.pythonhosted.org/packages/e1/32/df7f18bd0e724e0d9748829765455d6643ec847b3f87e77456fc99d0edab/black-24.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:fe4d6476887de70546212c99ac9bd803d90b42fc4767f058a0baa895013fbb3e", size = 1414209, upload-time = "2024-10-07T19:24:42.54Z" }, - { url = "https://files.pythonhosted.org/packages/c2/cc/7496bb63a9b06a954d3d0ac9fe7a73f3bf1cd92d7a58877c27f4ad1e9d41/black-24.10.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5a2221696a8224e335c28816a9d331a6c2ae15a2ee34ec857dcf3e45dbfa99ad", size = 1607468, upload-time = "2024-10-07T19:26:14.966Z" }, - { url = "https://files.pythonhosted.org/packages/2b/e3/69a738fb5ba18b5422f50b4f143544c664d7da40f09c13969b2fd52900e0/black-24.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f9da3333530dbcecc1be13e69c250ed8dfa67f43c4005fb537bb426e19200d50", size = 1437270, upload-time = "2024-10-07T19:25:24.291Z" }, - { url = "https://files.pythonhosted.org/packages/c9/9b/2db8045b45844665c720dcfe292fdaf2e49825810c0103e1191515fc101a/black-24.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4007b1393d902b48b36958a216c20c4482f601569d19ed1df294a496eb366392", size = 1737061, upload-time = "2024-10-07T19:23:52.18Z" }, - { url = "https://files.pythonhosted.org/packages/a3/95/17d4a09a5be5f8c65aa4a361444d95edc45def0de887810f508d3f65db7a/black-24.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:394d4ddc64782e51153eadcaaca95144ac4c35e27ef9b0a42e121ae7e57a9175", size = 1423293, upload-time = "2024-10-07T19:24:41.7Z" }, - { url = "https://files.pythonhosted.org/packages/90/04/bf74c71f592bcd761610bbf67e23e6a3cff824780761f536512437f1e655/black-24.10.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b5e39e0fae001df40f95bd8cc36b9165c5e2ea88900167bddf258bacef9bbdc3", size = 1644256, upload-time = "2024-10-07T19:27:53.355Z" }, - { url = "https://files.pythonhosted.org/packages/4c/ea/a77bab4cf1887f4b2e0bce5516ea0b3ff7d04ba96af21d65024629afedb6/black-24.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d37d422772111794b26757c5b55a3eade028aa3fde43121ab7b673d050949d65", size = 1448534, upload-time = "2024-10-07T19:26:44.953Z" }, - { url = "https://files.pythonhosted.org/packages/4e/3e/443ef8bc1fbda78e61f79157f303893f3fddf19ca3c8989b163eb3469a12/black-24.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14b3502784f09ce2443830e3133dacf2c0110d45191ed470ecb04d0f5f6fcb0f", size = 1761892, upload-time = "2024-10-07T19:24:10.264Z" }, - { url = "https://files.pythonhosted.org/packages/52/93/eac95ff229049a6901bc84fec6908a5124b8a0b7c26ea766b3b8a5debd22/black-24.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:30d2c30dc5139211dda799758559d1b049f7f14c580c409d6ad925b74a4208a8", size = 1434796, upload-time = "2024-10-07T19:25:06.239Z" }, - { url = "https://files.pythonhosted.org/packages/d0/a0/a993f58d4ecfba035e61fca4e9f64a2ecae838fc9f33ab798c62173ed75c/black-24.10.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1cbacacb19e922a1d75ef2b6ccaefcd6e93a2c05ede32f06a21386a04cedb981", size = 1643986, upload-time = "2024-10-07T19:28:50.684Z" }, - { url = "https://files.pythonhosted.org/packages/37/d5/602d0ef5dfcace3fb4f79c436762f130abd9ee8d950fa2abdbf8bbc555e0/black-24.10.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1f93102e0c5bb3907451063e08b9876dbeac810e7da5a8bfb7aeb5a9ef89066b", size = 1448085, upload-time = "2024-10-07T19:28:12.093Z" }, - { url = "https://files.pythonhosted.org/packages/47/6d/a3a239e938960df1a662b93d6230d4f3e9b4a22982d060fc38c42f45a56b/black-24.10.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ddacb691cdcdf77b96f549cf9591701d8db36b2f19519373d60d31746068dbf2", size = 1760928, upload-time = "2024-10-07T19:24:15.233Z" }, - { url = "https://files.pythonhosted.org/packages/dd/cf/af018e13b0eddfb434df4d9cd1b2b7892bab119f7a20123e93f6910982e8/black-24.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:680359d932801c76d2e9c9068d05c6b107f2584b2a5b88831c83962eb9984c1b", size = 1436875, upload-time = "2024-10-07T19:24:42.762Z" }, - { url = "https://files.pythonhosted.org/packages/8d/a7/4b27c50537ebca8bec139b872861f9d2bf501c5ec51fcf897cb924d9e264/black-24.10.0-py3-none-any.whl", hash = "sha256:3bb2b7a1f7b685f85b11fed1ef10f8a9148bceb49853e47a294a3dd963c1dd7d", size = 206898, upload-time = "2024-10-07T19:20:48.317Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/8c/ad/33adf4708633d047950ff2dfdea2e215d84ac50ef95aff14a614e4b6e9b2/black-25.11.0.tar.gz", hash = "sha256:9a323ac32f5dc75ce7470501b887250be5005a01602e931a15e45593f70f6e08", size = 655669, upload-time = "2025-11-10T01:53:50.558Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/d2/6caccbc96f9311e8ec3378c296d4f4809429c43a6cd2394e3c390e86816d/black-25.11.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ec311e22458eec32a807f029b2646f661e6859c3f61bc6d9ffb67958779f392e", size = 1743501, upload-time = "2025-11-10T01:59:06.202Z" }, + { url = "https://files.pythonhosted.org/packages/69/35/b986d57828b3f3dccbf922e2864223197ba32e74c5004264b1c62bc9f04d/black-25.11.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1032639c90208c15711334d681de2e24821af0575573db2810b0763bcd62e0f0", size = 1597308, upload-time = "2025-11-10T01:57:58.633Z" }, + { url = "https://files.pythonhosted.org/packages/39/8e/8b58ef4b37073f52b64a7b2dd8c9a96c84f45d6f47d878d0aa557e9a2d35/black-25.11.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0c0f7c461df55cf32929b002335883946a4893d759f2df343389c4396f3b6b37", size = 1656194, upload-time = "2025-11-10T01:57:10.909Z" }, + { url = "https://files.pythonhosted.org/packages/8d/30/9c2267a7955ecc545306534ab88923769a979ac20a27cf618d370091e5dd/black-25.11.0-cp310-cp310-win_amd64.whl", hash = "sha256:f9786c24d8e9bd5f20dc7a7f0cdd742644656987f6ea6947629306f937726c03", size = 1347996, upload-time = "2025-11-10T01:57:22.391Z" }, + { url = "https://files.pythonhosted.org/packages/c4/62/d304786b75ab0c530b833a89ce7d997924579fb7484ecd9266394903e394/black-25.11.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:895571922a35434a9d8ca67ef926da6bc9ad464522a5fe0db99b394ef1c0675a", size = 1727891, upload-time = "2025-11-10T02:01:40.507Z" }, + { url = "https://files.pythonhosted.org/packages/82/5d/ffe8a006aa522c9e3f430e7b93568a7b2163f4b3f16e8feb6d8c3552761a/black-25.11.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:cb4f4b65d717062191bdec8e4a442539a8ea065e6af1c4f4d36f0cdb5f71e170", size = 1581875, upload-time = "2025-11-10T01:57:51.192Z" }, + { url = "https://files.pythonhosted.org/packages/cb/c8/7c8bda3108d0bb57387ac41b4abb5c08782b26da9f9c4421ef6694dac01a/black-25.11.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d81a44cbc7e4f73a9d6ae449ec2317ad81512d1e7dce7d57f6333fd6259737bc", size = 1642716, upload-time = "2025-11-10T01:56:51.589Z" }, + { url = "https://files.pythonhosted.org/packages/34/b9/f17dea34eecb7cc2609a89627d480fb6caea7b86190708eaa7eb15ed25e7/black-25.11.0-cp311-cp311-win_amd64.whl", hash = "sha256:7eebd4744dfe92ef1ee349dc532defbf012a88b087bb7ddd688ff59a447b080e", size = 1352904, upload-time = "2025-11-10T01:59:26.252Z" }, + { url = "https://files.pythonhosted.org/packages/7f/12/5c35e600b515f35ffd737da7febdb2ab66bb8c24d88560d5e3ef3d28c3fd/black-25.11.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:80e7486ad3535636657aa180ad32a7d67d7c273a80e12f1b4bfa0823d54e8fac", size = 1772831, upload-time = "2025-11-10T02:03:47Z" }, + { url = "https://files.pythonhosted.org/packages/1a/75/b3896bec5a2bb9ed2f989a970ea40e7062f8936f95425879bbe162746fe5/black-25.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6cced12b747c4c76bc09b4db057c319d8545307266f41aaee665540bc0e04e96", size = 1608520, upload-time = "2025-11-10T01:58:46.895Z" }, + { url = "https://files.pythonhosted.org/packages/f3/b5/2bfc18330eddbcfb5aab8d2d720663cd410f51b2ed01375f5be3751595b0/black-25.11.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6cb2d54a39e0ef021d6c5eef442e10fd71fcb491be6413d083a320ee768329dd", size = 1682719, upload-time = "2025-11-10T01:56:55.24Z" }, + { url = "https://files.pythonhosted.org/packages/96/fb/f7dc2793a22cdf74a72114b5ed77fe3349a2e09ef34565857a2f917abdf2/black-25.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:ae263af2f496940438e5be1a0c1020e13b09154f3af4df0835ea7f9fe7bfa409", size = 1362684, upload-time = "2025-11-10T01:57:07.639Z" }, + { url = "https://files.pythonhosted.org/packages/ad/47/3378d6a2ddefe18553d1115e36aea98f4a90de53b6a3017ed861ba1bd3bc/black-25.11.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0a1d40348b6621cc20d3d7530a5b8d67e9714906dfd7346338249ad9c6cedf2b", size = 1772446, upload-time = "2025-11-10T02:02:16.181Z" }, + { url = "https://files.pythonhosted.org/packages/ba/4b/0f00bfb3d1f7e05e25bfc7c363f54dc523bb6ba502f98f4ad3acf01ab2e4/black-25.11.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:51c65d7d60bb25429ea2bf0731c32b2a2442eb4bd3b2afcb47830f0b13e58bfd", size = 1607983, upload-time = "2025-11-10T02:02:52.502Z" }, + { url = "https://files.pythonhosted.org/packages/99/fe/49b0768f8c9ae57eb74cc10a1f87b4c70453551d8ad498959721cc345cb7/black-25.11.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:936c4dd07669269f40b497440159a221ee435e3fddcf668e0c05244a9be71993", size = 1682481, upload-time = "2025-11-10T01:57:12.35Z" }, + { url = "https://files.pythonhosted.org/packages/55/17/7e10ff1267bfa950cc16f0a411d457cdff79678fbb77a6c73b73a5317904/black-25.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:f42c0ea7f59994490f4dccd64e6b2dd49ac57c7c84f38b8faab50f8759db245c", size = 1363869, upload-time = "2025-11-10T01:58:24.608Z" }, + { url = "https://files.pythonhosted.org/packages/67/c0/cc865ce594d09e4cd4dfca5e11994ebb51604328489f3ca3ae7bb38a7db5/black-25.11.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:35690a383f22dd3e468c85dc4b915217f87667ad9cce781d7b42678ce63c4170", size = 1771358, upload-time = "2025-11-10T02:03:33.331Z" }, + { url = "https://files.pythonhosted.org/packages/37/77/4297114d9e2fd2fc8ab0ab87192643cd49409eb059e2940391e7d2340e57/black-25.11.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:dae49ef7369c6caa1a1833fd5efb7c3024bb7e4499bf64833f65ad27791b1545", size = 1612902, upload-time = "2025-11-10T01:59:33.382Z" }, + { url = "https://files.pythonhosted.org/packages/de/63/d45ef97ada84111e330b2b2d45e1dd163e90bd116f00ac55927fb6bf8adb/black-25.11.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5bd4a22a0b37401c8e492e994bce79e614f91b14d9ea911f44f36e262195fdda", size = 1680571, upload-time = "2025-11-10T01:57:04.239Z" }, + { url = "https://files.pythonhosted.org/packages/ff/4b/5604710d61cdff613584028b4cb4607e56e148801ed9b38ee7970799dab6/black-25.11.0-cp314-cp314-win_amd64.whl", hash = "sha256:aa211411e94fdf86519996b7f5f05e71ba34835d8f0c0f03c00a26271da02664", size = 1382599, upload-time = "2025-11-10T01:57:57.427Z" }, + { url = "https://files.pythonhosted.org/packages/00/5d/aed32636ed30a6e7f9efd6ad14e2a0b0d687ae7c8c7ec4e4a557174b895c/black-25.11.0-py3-none-any.whl", hash = "sha256:e3f562da087791e96cefcd9dda058380a442ab322a02e222add53736451f604b", size = 204918, upload-time = "2025-11-10T01:53:48.917Z" }, ] [[package]] @@ -1393,8 +1398,8 @@ wheels = [ [[package]] name = "gt4py" -version = "1.1.1+unknown.version.details" -source = { git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion#9e75186efddebb1eba0003075688e140a2b7ba75" } +version = "1.1.1.post18+f80482dc" +source = { git = "https://github.com/GridTools/gt4py?branch=main#f80482dc45fe2b8e58de84be39961a8fa48fdce2" } dependencies = [ { name = "attrs" }, { name = "black" }, @@ -1740,7 +1745,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", git = "https://github.com/GridTools/gt4py?branch=main" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1757,7 +1762,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", git = "https://github.com/GridTools/gt4py?branch=main" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1774,7 +1779,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", git = "https://github.com/GridTools/gt4py?branch=main" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1791,7 +1796,7 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", git = "https://github.com/GridTools/gt4py?branch=main" }, { name = "icon4py-common", editable = "model/common" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1803,13 +1808,15 @@ source = { editable = "model/atmosphere/subgrid_scale_physics/muphys" } dependencies = [ { name = "gt4py" }, { name = "icon4py-common", extra = ["io"] }, + { name = "numpy" }, { name = "packaging" }, ] [package.metadata] requires-dist = [ - { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", git = "https://github.com/GridTools/gt4py?branch=main" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, + { name = "numpy", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, ] @@ -1873,9 +1880,9 @@ requires-dist = [ { name = "dace", git = "https://github.com/GridTools/dace?tag=__gt4py-next-integration_2025_11_05" }, { name = "datashader", marker = "extra == 'io'", specifier = ">=0.16.1" }, { name = "ghex", marker = "extra == 'distributed'", git = "https://github.com/msimberg/GHEX.git?branch=async-mpi" }, - { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, - { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, - { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", git = "https://github.com/GridTools/gt4py?branch=main" }, + { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/GridTools/gt4py?branch=main" }, + { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/GridTools/gt4py?branch=main" }, { name = "holoviews", marker = "extra == 'io'", specifier = ">=1.16.0" }, { name = "icon4py-common", extras = ["distributed", "io"], marker = "extra == 'all'", editable = "model/common" }, { name = "mpi4py", marker = "extra == 'distributed'", specifier = ">=3.1.5" }, @@ -1911,7 +1918,7 @@ dependencies = [ requires-dist = [ { name = "click", specifier = ">=8.0.1" }, { name = "devtools", specifier = ">=0.12" }, - { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", git = "https://github.com/GridTools/gt4py?branch=main" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, { name = "icon4py-common", editable = "model/common" }, @@ -1939,7 +1946,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "filelock", specifier = ">=3.18.0" }, - { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", git = "https://github.com/GridTools/gt4py?branch=main" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, { name = "numpy", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, @@ -1988,9 +1995,9 @@ requires-dist = [ { name = "cupy-cuda11x", marker = "extra == 'cuda11'", specifier = ">=13.0" }, { name = "cupy-cuda12x", marker = "extra == 'cuda12'", specifier = ">=13.0" }, { name = "fprettify", specifier = ">=0.3.7" }, - { name = "gt4py", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, - { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, - { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/havogt/gt4py?branch=cse_in_fieldop_fusion" }, + { name = "gt4py", git = "https://github.com/GridTools/gt4py?branch=main" }, + { name = "gt4py", extras = ["cuda11"], marker = "extra == 'cuda11'", git = "https://github.com/GridTools/gt4py?branch=main" }, + { name = "gt4py", extras = ["cuda12"], marker = "extra == 'cuda12'", git = "https://github.com/GridTools/gt4py?branch=main" }, { name = "icon4py-atmosphere-advection", editable = "model/atmosphere/advection" }, { name = "icon4py-atmosphere-diffusion", editable = "model/atmosphere/diffusion" }, { name = "icon4py-atmosphere-dycore", editable = "model/atmosphere/dycore" }, @@ -3587,6 +3594,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6a/3e/b68c118422ec867fa7ab88444e1274aa40681c606d59ac27de5a5588f082/python_dotenv-1.0.1-py3-none-any.whl", hash = "sha256:f7b63ef50f1b690dddf550d03497b66d609393b40b564ed0d674909a68ebf16a", size = 19863, upload-time = "2024-01-23T06:32:58.246Z" }, ] +[[package]] +name = "pytokens" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/8d/a762be14dae1c3bf280202ba3172020b2b0b4c537f94427435f19c413b72/pytokens-0.3.0.tar.gz", hash = "sha256:2f932b14ed08de5fcf0b391ace2642f858f1394c0857202959000b68ed7a458a", size = 17644, upload-time = "2025-11-05T13:36:35.34Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/25/d9db8be44e205a124f6c98bc0324b2bb149b7431c53877fc6d1038dddaf5/pytokens-0.3.0-py3-none-any.whl", hash = "sha256:95b2b5eaf832e469d141a378872480ede3f251a5a5041b8ec6e581d3ac71bbf3", size = 12195, upload-time = "2025-11-05T13:36:33.183Z" }, +] + [[package]] name = "pytz" version = "2024.2" From 3ff602f5a9f6baf9f29634d92197acbb17d43dbd Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Fri, 5 Dec 2025 17:28:43 +0100 Subject: [PATCH 21/23] fix module import --- .../src/icon4py/model/common/orchestration/decorator.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model/common/src/icon4py/model/common/orchestration/decorator.py b/model/common/src/icon4py/model/common/orchestration/decorator.py index 17ff30e0de..834ef47270 100644 --- a/model/common/src/icon4py/model/common/orchestration/decorator.py +++ b/model/common/src/icon4py/model/common/orchestration/decorator.py @@ -31,7 +31,7 @@ from dace import hooks from dace.transformation.passes import simplify as dace_simplify from gt4py._core import definitions as core_defs -from gt4py.next.program_processors.runners.dace import utils as gtx_dace_utils +from gt4py.next.program_processors.runners.dace import sdfg_args as gtx_dace_args from icon4py.model.common import dimension as dims from icon4py.model.common.decomposition import definitions as decomposition @@ -159,7 +159,7 @@ def wrapper(*args, **kwargs): { k: v for k, v in grid.connectivities.items() - if gtx_dace_utils.connectivity_identifier(k) in sdfg.arrays + if gtx_dace_args.connectivity_identifier(k) in sdfg.arrays }, ), } @@ -509,7 +509,7 @@ def dace_specific_kwargs( return { # connectivity tables at runtime **{ - gtx_dace_utils.connectivity_identifier(k): v.ndarray + gtx_dace_args.connectivity_identifier(k): v.ndarray for k, v in offset_providers.items() if hasattr(v, "ndarray") }, From 4838d26776420775dce483a86611053e152dd50e Mon Sep 17 00:00:00 2001 From: Edoardo Paone Date: Tue, 9 Dec 2025 18:09:53 +0100 Subject: [PATCH 22/23] undo extra changes --- model/testing/src/icon4py/model/testing/stencil_tests.py | 5 +---- uv.lock | 2 -- 2 files changed, 1 insertion(+), 6 deletions(-) diff --git a/model/testing/src/icon4py/model/testing/stencil_tests.py b/model/testing/src/icon4py/model/testing/stencil_tests.py index 92646cf12f..f83798f029 100644 --- a/model/testing/src/icon4py/model/testing/stencil_tests.py +++ b/model/testing/src/icon4py/model/testing/stencil_tests.py @@ -38,10 +38,7 @@ def allocate_data( ) -> dict[str, gtx.Field | tuple[gtx.Field, ...]]: _allocate_field = constructors.as_field.partial(allocator=allocator) # type:ignore[attr-defined] # TODO(havogt): check why it doesn't understand the fluid_partial input_data = { - # TODO: make this work for dataclass named collections as well - k: v.__class__( - *tuple(_allocate_field(domain=field.domain, data=field.ndarray) for field in v) - ) + k: tuple(_allocate_field(domain=field.domain, data=field.ndarray) for field in v) if isinstance(v, tuple) else _allocate_field(domain=v.domain, data=v.ndarray) if not gtx.is_scalar_type(v) and k != "domain" diff --git a/uv.lock b/uv.lock index fdf72af276..c6539f9b4d 100644 --- a/uv.lock +++ b/uv.lock @@ -1812,7 +1812,6 @@ source = { editable = "model/atmosphere/subgrid_scale_physics/muphys" } dependencies = [ { name = "gt4py" }, { name = "icon4py-common", extra = ["io"] }, - { name = "numpy" }, { name = "packaging" }, ] @@ -1820,7 +1819,6 @@ dependencies = [ requires-dist = [ { name = "gt4py", specifier = "==1.1.2" }, { name = "icon4py-common", extras = ["io"], editable = "model/common" }, - { name = "numpy", specifier = ">=1.23.3" }, { name = "packaging", specifier = ">=20.0" }, ] From 42e4b86782d2ef4daaa56c59608ad9590940a78e Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Fri, 9 Jan 2026 09:06:00 +0100 Subject: [PATCH 23/23] fix timers --- .../subgrid_scale_physics/muphys/driver/run_full_muphys.py | 4 ++-- .../subgrid_scale_physics/muphys/driver/run_graupel_only.py | 4 ++-- model/common/src/icon4py/model/common/utils/device_utils.py | 1 + 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py index f67446d810..75c3c058f3 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_full_muphys.py @@ -184,7 +184,7 @@ def main(): start_time = None for _x in range(int(args.itime) + 1): if _x == 1: # Only start timing second iteration - device_utils.sync(backend) + device_utils.sync(allocator) start_time = time.time() muphys_step( @@ -203,7 +203,7 @@ def main(): pre=out.pre, ) - device_utils.sync(backend) + device_utils.sync(allocator) end_time = time.time() if start_time is not None: diff --git a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py index 865397d86b..6a1d3bb560 100755 --- a/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py +++ b/model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py @@ -78,7 +78,7 @@ def main(): start_time = None for _x in range(int(args.itime) + 1): if _x == 1: # Only start timing second iteration - device_utils.sync(backend) + device_utils.sync(allocator) start_time = time.time() graupel_run_program( @@ -96,7 +96,7 @@ def main(): pg=out.pg, pre=out.pre, ) - device_utils.sync(backend) + device_utils.sync(allocator) end_time = time.time() if start_time is not None: diff --git a/model/common/src/icon4py/model/common/utils/device_utils.py b/model/common/src/icon4py/model/common/utils/device_utils.py index cacfc8eb64..0429096f70 100644 --- a/model/common/src/icon4py/model/common/utils/device_utils.py +++ b/model/common/src/icon4py/model/common/utils/device_utils.py @@ -37,6 +37,7 @@ def sync(allocator: gtx_typing.FieldBufferAllocationUtil | None = None) -> None: Note: this is and ad-hoc interface, maybe the function should get the device to sync for. """ + assert allocator is None or gtx_allocators.is_field_allocation_tool(allocator) if allocator is not None and is_cupy_device(allocator): cp.cuda.runtime.deviceSynchronize()