muphys: Only write the surface level of the fluxes

ci/dace.yml

@@ -19,7 +19,7 @@ include:
         GRID: [simple, icon_regional]
 # test_model_stencils_x86_64:
 #   extends: [.test_model_stencils, .test_template_x86_64]
-test_model_stencils_aarch64:
+.test_model_stencils_aarch64:
   extends: [.test_model_stencils, .test_template_aarch64]
 
 .test_model_datatests:
@@ -28,7 +28,7 @@ test_model_stencils_aarch64:
     - nox -s "test_model-3.10(datatest, $COMPONENT)" -- --backend=$BACKEND --level=$LEVEL
   parallel:
     matrix:
-      - COMPONENT: [advection, diffusion, dycore, microphysics, muphys, common, driver]
+      - COMPONENT: [muphys]
         BACKEND: [dace_cpu, dace_gpu]
         LEVEL: [integration]
   rules:
@@ -44,7 +44,7 @@ test_model_stencils_aarch64:
         SLURM_TIMELIMIT: '01:00:00'
     - when: on_success
       variables:
-        SLURM_TIMELIMIT: '00:30:00'
+        SLURM_TIMELIMIT: '01:00:00'
 # test_model_datatests_x86_64:
 #   extends: [.test_model_datatests, .test_template_x86_64]
 test_model_datatests_aarch64:

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/common.py

@@ -11,6 +11,7 @@
 import dataclasses
 import functools
 import pathlib
+from typing import ClassVar
 
 import netCDF4
 import numpy as np
@@ -143,6 +144,8 @@ class GraupelOutput:
     pg: gtx.Field[dims.CellDim, dims.KDim] | None
     pre: gtx.Field[dims.CellDim, dims.KDim] | None
 
+    _surface_fields: ClassVar[list[str]] = ["pr", "ps", "pi", "pg", "pre"]
+
     @classmethod
     def allocate(
         cls,
@@ -160,10 +163,22 @@ def allocate(
         if references is None:
             references = {}
 
-        zeros = functools.partial(gtx.zeros, domain=domain, allocator=allocator)
+        zeros_full = functools.partial(gtx.zeros, domain=domain, allocator=allocator)
+        surface_domain = gtx.Domain(
+            dims=domain.dims,
+            ranges=(
+                domain.ranges[0],
+                gtx.unit_range((domain.ranges[1].stop - 1, domain.ranges[1].stop)),
+            ),
+        )
+        zeros_surface = functools.partial(gtx.zeros, domain=surface_domain, allocator=allocator)
         return cls(
             **{
-                field.name: zeros() if field.name not in references else references[field.name]
+                field.name: references[field.name]
+                if field.name in references
+                else zeros_surface()
+                if field.name in cls._surface_fields
+                else zeros_full()
                 for field in dataclasses.fields(cls)
             }
         )
@@ -210,6 +225,9 @@ def write(self, filename: pathlib.Path | str):
             write_height_field = functools.partial(
                 _field_to_nc, ncfile, ("height", "ncells"), dtype=np.float64
             )
+            write_surface_field = functools.partial(
+                _field_to_nc, ncfile, ("surface", "ncells"), dtype=np.float64
+            )
 
             write_height_field("ta", self.t)
             write_height_field("hus", self.qv)
@@ -221,14 +239,12 @@ def write(self, filename: pathlib.Path | str):
             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(havogt): see https://github.com/C2SM/icon4py/pull/995
+                write_surface_field("prr_gsp", self.pr)
             if self.ps is not None:
-                write_height_field("prs_gsp", self.ps)  # TODO(havogt): see above
+                write_surface_field("prs_gsp", self.ps)
             if self.pi is not None:
-                write_height_field("pri_gsp", self.pi)  # TODO(havogt): see above
+                write_surface_field("pri_gsp", self.pi)
             if self.pg is not None:
-                write_height_field("prg_gsp", self.pg)  # TODO(havogt): see above
+                write_surface_field("prg_gsp", self.pg)
             if self.pre is not None:
-                write_height_field("pre_gsp", self.pre)  # TODO(havogt): see above
+                write_surface_field("pre_gsp", self.pre)

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/driver/run_graupel_only.py

@@ -88,7 +88,7 @@ def main():
     use_inout_buffers = True  # Set to True to reuse input buffers for output.
     if use_inout_buffers:
         # We are passing the same buffers for `Q` as input and output. This is not best GT4Py practice,
-        # but should be save in this case as we are not reading the input with an offset.
+        # but should be safe in this case as we are not reading the input with an offset.
         references = {
             "qv": inp.qv,
             "qc": inp.qc,

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/graupel.py

@@ -594,8 +594,68 @@ def graupel_run(
         qnc=qnc,
         enable_masking=enable_masking,
         out=(t_out, q_out, pflx, pr, ps, pi, pg, pre),
-        domain={
-            dims.CellDim: (horizontal_start, horizontal_end),
-            dims.KDim: (vertical_start, vertical_end),
-        },
+        domain=(
+            # t_out
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_start, vertical_end),
+            },
+            # q_out
+            (
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+            ),
+            # pflx
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_start, vertical_end),
+            },
+            # pr
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+            # ps
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+            # pi
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+            # pg
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+            # pre
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+        ),
     )

model/atmosphere/subgrid_scale_physics/muphys/src/icon4py/model/atmosphere/subgrid_scale_physics/muphys/implementations/muphys.py

@@ -95,8 +95,68 @@ def muphys_run(
         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),
-        },
+        domain=(
+            # t_out
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_start, vertical_end),
+            },
+            # q_out
+            (
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+                {
+                    dims.CellDim: (horizontal_start, horizontal_end),
+                    dims.KDim: (vertical_start, vertical_end),
+                },
+            ),
+            # pflx
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_start, vertical_end),
+            },
+            # pr
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+            # ps
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+            # pi
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+            # pg
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+            # pre
+            {
+                dims.CellDim: (horizontal_start, horizontal_end),
+                dims.KDim: (vertical_end - 1, vertical_end),
+            },
+        ),
     )

pyproject.toml

@@ -362,7 +362,7 @@ url = 'https://gridtools.github.io/pypi/'
 [tool.uv.sources]
 dace = {index = "gridtools"}
 ghex = {git = "https://github.com/msimberg/GHEX.git", branch = "async-mpi"}
-# gt4py = {git = "https://github.com/GridTools/gt4py", branch = "main"}
+gt4py = {git = "https://github.com/edopao/gt4py", branch = "dace_scan_single_level"}
 # gt4py = {index = "test.pypi"}
 icon4py-atmosphere-advection = {workspace = true}
 icon4py-atmosphere-diffusion = {workspace = true}