Add advection granule with horizontal transport

ci/default.yml

@@ -34,7 +34,7 @@ test_tools_aarch64:
     - tox -r -e run_model_tests -c model/ --verbose -- --backend=$BACKEND $COMPONENT
   parallel:
     matrix:
-    - COMPONENT: [atmosphere/diffusion/tests/diffusion_tests, atmosphere/dycore/tests/dycore_tests, atmosphere/subgrid_scale_physics/microphysics/tests, common/tests, driver/tests]
+    - COMPONENT: [atmosphere/advection/tests/advection_tests atmosphere/diffusion/tests/diffusion_tests, atmosphere/dycore/tests/dycore_tests, atmosphere/subgrid_scale_physics/microphysics/tests, common/tests, driver/tests]
       BACKEND: [gtfn_cpu]
 test_model_datatests_x86_64:
   extends: [.test_model_datatests, .test_template_x86_64]

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/advection_states.py

@@ -0,0 +1,75 @@
+# 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 dataclasses
+
+import gt4py.next as gtx
+
+from icon4py.model.common import dimension as dims, field_type_aliases as fa
+from icon4py.model.common.type_alias import wpfloat
+
+
+@dataclasses.dataclass(frozen=True)
+class AdvectionDiagnosticState:
+    """Represents the diagnostic fields needed in advection."""
+
+    # mass of air in layer at physics time step now [kg/m^2] (nproma,nlev,nblks_c)
+    airmass_now: fa.CellKField[wpfloat]
+
+    # mass of air in layer at physics time step new [kg/m^2] (nproma,nlev,nblks_c)
+    airmass_new: fa.CellKField[wpfloat]
+
+    # tracer tendency field for use in grid refinement [kg/kg/s] (nproma,nlev,nblks_c)
+    grf_tend_tracer: fa.CellKField[wpfloat]
+
+    # horizontal tracer flux at edges [kg/m/s] (nproma,nlev,nblks_e)
+    hfl_tracer: fa.EdgeKField[wpfloat]
+
+    # vertical tracer flux at cells [kg/m/s] (nproma,nlevp1,nblks_c)
+    # TODO (dastrm): should be KHalfDim
+    vfl_tracer: fa.CellKField[wpfloat]
+
+
+@dataclasses.dataclass(frozen=True)
+class AdvectionInterpolationState:
+    """Represents the interpolation state needed in advection."""
+
+    # factor for divergence (nproma,cell_type,nblks_c)
+    geofac_div: gtx.Field[[dims.CEDim], wpfloat]
+
+    # coefficients used for rbf interpolation of the tangential velocity component (rbf_vec_dim_e,nproma,nblks_e)
+    rbf_vec_coeff_e: gtx.Field[[dims.EdgeDim, dims.E2C2EDim], wpfloat]
+
+    # x-components of positions of various points on local plane tangential to the edge midpoint (nproma,4,nblks_e)
+    pos_on_tplane_e_1: gtx.Field[[dims.ECDim], wpfloat]
+
+    # y-components of positions of various points on local plane tangential to the edge midpoint (nproma,4,nblks_e)
+    pos_on_tplane_e_2: gtx.Field[[dims.ECDim], wpfloat]
+
+
+@dataclasses.dataclass(frozen=True)
+class AdvectionLeastSquaresState:
+    """Represents the least squares state needed in advection."""
+
+    # pseudo (or Moore-Penrose) inverse of lsq design matrix A, originally (nproma,lsq_dim_unk,lsq_dim_c,nblks_c)
+    lsq_pseudoinv_1: gtx.Field[[dims.CECDim], wpfloat]
+    lsq_pseudoinv_2: gtx.Field[[dims.CECDim], wpfloat]
+
+
+@dataclasses.dataclass(frozen=True)
+class AdvectionMetricState:
+    """Represents the metric fields needed in advection."""
+
+    # metrical modification factor for horizontal part of divergence at full levels (nlev)
+    deepatmo_divh: fa.KField[wpfloat]
+
+    # metrical modification factor for vertical part of divergence at full levels (nlev)
+    deepatmo_divzl: fa.KField[wpfloat]
+
+    # metrical modification factor for vertical part of divergence at full levels (nlev)
+    deepatmo_divzu: fa.KField[wpfloat]

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/__init__.py

@@ -0,0 +1,8 @@
+# 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
+

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_02.py → model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/apply_horizontal_density_increment.py

@@ -8,14 +8,14 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import maximum
+from gt4py.next.ffront.fbuiltins import int32, maximum
 
 from icon4py.model.common import field_type_aliases as fa
-from icon4py.model.common.dimension import Koff
+from icon4py.model.common.dimension import CellDim, KDim, Koff
 
 
 @field_operator
-def _step_advection_stencil_02(
+def _apply_horizontal_density_increment(
     p_rhodz_new: fa.CellKField[float],
     p_mflx_contra_v: fa.CellKField[float],
     deepatmo_divzl: fa.KField[float],
@@ -28,19 +28,24 @@ def _step_advection_stencil_02(
 
 
 @program(grid_type=GridType.UNSTRUCTURED)
-def step_advection_stencil_02(
+def apply_horizontal_density_increment(
     p_rhodz_new: fa.CellKField[float],
     p_mflx_contra_v: fa.CellKField[float],
     deepatmo_divzl: fa.KField[float],
     deepatmo_divzu: fa.KField[float],
     p_dtime: float,
     rhodz_ast2: fa.CellKField[float],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
-    _step_advection_stencil_02(
+    _apply_horizontal_density_increment(
         p_rhodz_new,
         p_mflx_contra_v,
         deepatmo_divzl,
         deepatmo_divzu,
         p_dtime,
         out=rhodz_ast2,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
     )

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_03.py → model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/apply_interpolated_tracer_time_tendency.py

@@ -8,13 +8,14 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import maximum
+from gt4py.next.ffront.fbuiltins import int32, maximum
 
 from icon4py.model.common import field_type_aliases as fa
+from icon4py.model.common.dimension import CellDim, KDim
 
 
 @field_operator
-def _step_advection_stencil_03(
+def _apply_interpolated_tracer_time_tendency(
     p_tracer_now: fa.CellKField[float],
     p_grf_tend_tracer: fa.CellKField[float],
     p_dtime: float,
@@ -24,10 +25,20 @@ def _step_advection_stencil_03(
 
 
 @program(grid_type=GridType.UNSTRUCTURED)
-def step_advection_stencil_03(
+def apply_interpolated_tracer_time_tendency(
     p_tracer_now: fa.CellKField[float],
     p_grf_tend_tracer: fa.CellKField[float],
     p_tracer_new: fa.CellKField[float],
     p_dtime: float,
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
-    _step_advection_stencil_03(p_tracer_now, p_grf_tend_tracer, p_dtime, out=p_tracer_new)
+    _apply_interpolated_tracer_time_tendency(
+        p_tracer_now,
+        p_grf_tend_tracer,
+        p_dtime,
+        out=p_tracer_new,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
+    )

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_mo_stencil_04.py → model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/apply_monotone_horizontal_multiplicative_flux_factors.py

@@ -14,7 +14,7 @@
 
 
 @field_operator
-def _hflx_limiter_mo_stencil_04(
+def _apply_monotone_horizontal_multiplicative_flux_factors(
     z_anti: fa.EdgeKField[float],
     r_m: fa.CellKField[float],
     r_p: fa.CellKField[float],
@@ -29,11 +29,13 @@ def _hflx_limiter_mo_stencil_04(
 
 
 @program
-def hflx_limiter_mo_stencil_04(
+def apply_monotone_horizontal_multiplicative_flux_factors(
     z_anti: fa.EdgeKField[float],
     r_m: fa.CellKField[float],
     r_p: fa.CellKField[float],
     z_mflx_low: fa.EdgeKField[float],
     p_mflx_tracer_h: fa.EdgeKField[float],
 ):
-    _hflx_limiter_mo_stencil_04(z_anti, r_m, r_p, z_mflx_low, out=p_mflx_tracer_h)
+    _apply_monotone_horizontal_multiplicative_flux_factors(
+        z_anti, r_m, r_p, z_mflx_low, out=p_mflx_tracer_h
+    )

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/hflx_limiter_pd_stencil_02.py → model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/apply_positive_definite_horizontal_multiplicative_flux_factor.py

@@ -7,14 +7,14 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import where
+from gt4py.next.ffront.fbuiltins import int32, where
 
 from icon4py.model.common import field_type_aliases as fa
-from icon4py.model.common.dimension import E2C
+from icon4py.model.common.dimension import E2C, EdgeDim, KDim
 
 
 @field_operator
-def _hflx_limiter_pd_stencil_02(
+def _apply_positive_definite_horizontal_multiplicative_flux_factor(
     r_m: fa.CellKField[float],
     p_mflx_tracer_h: fa.EdgeKField[float],
 ) -> fa.EdgeKField[float]:
@@ -27,12 +27,17 @@ def _hflx_limiter_pd_stencil_02(
 
 
 @program
-def hflx_limiter_pd_stencil_02(
+def apply_positive_definite_horizontal_multiplicative_flux_factor(
     r_m: fa.CellKField[float],
     p_mflx_tracer_h: fa.EdgeKField[float],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
-    _hflx_limiter_pd_stencil_02(
+    _apply_positive_definite_horizontal_multiplicative_flux_factor(
         r_m,
         p_mflx_tracer_h,
         out=p_mflx_tracer_h,
+        domain={EdgeDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
     )

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/step_advection_stencil_01.py → model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/apply_vertical_density_increment.py

@@ -8,13 +8,14 @@
 
 from gt4py.next.common import GridType
 from gt4py.next.ffront.decorator import field_operator, program
+from gt4py.next.ffront.fbuiltins import int32
 
 from icon4py.model.common import field_type_aliases as fa
-from icon4py.model.common.dimension import Koff
+from icon4py.model.common.dimension import CellDim, KDim, Koff
 
 
 @field_operator
-def _step_advection_stencil_01(
+def _apply_vertical_density_increment(
     rhodz_ast: fa.CellKField[float],
     p_mflx_contra_v: fa.CellKField[float],
     deepatmo_divzl: fa.KField[float],
@@ -28,19 +29,24 @@ def _step_advection_stencil_01(
 
 
 @program(grid_type=GridType.UNSTRUCTURED)
-def step_advection_stencil_01(
+def apply_vertical_density_increment(
     rhodz_ast: fa.CellKField[float],
     p_mflx_contra_v: fa.CellKField[float],
     deepatmo_divzl: fa.KField[float],
     deepatmo_divzu: fa.KField[float],
     p_dtime: float,
     rhodz_ast2: fa.CellKField[float],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
 ):
-    _step_advection_stencil_01(
+    _apply_vertical_density_increment(
         rhodz_ast,
         p_mflx_contra_v,
         deepatmo_divzl,
         deepatmo_divzu,
         p_dtime,
         out=rhodz_ast2,
+        domain={CellDim: (horizontal_start, horizontal_end), KDim: (vertical_start, vertical_end)},
     )

model/atmosphere/advection/src/icon4py/model/atmosphere/advection/upwind_hflux_miura_cycl_stencil_03a.py → model/atmosphere/advection/src/icon4py/model/atmosphere/advection/stencils/average_horizontal_flux_subcycling_2.py

@@ -14,7 +14,7 @@
 
 
 @field_operator
-def _upwind_hflux_miura_cycl_stencil_03a(
+def _average_horizontal_flux_subcycling_2(
     z_tracer_mflx_1_dsl: fa.EdgeKField[float],
     z_tracer_mflx_2_dsl: fa.EdgeKField[float],
 ) -> fa.EdgeKField[float]:
@@ -23,9 +23,9 @@ def _upwind_hflux_miura_cycl_stencil_03a(
 
 
 @program(grid_type=GridType.UNSTRUCTURED)
-def upwind_hflux_miura_cycl_stencil_03a(
+def average_horizontal_flux_subcycling_2(
     z_tracer_mflx_1_dsl: fa.EdgeKField[float],
     z_tracer_mflx_2_dsl: fa.EdgeKField[float],
     p_out_e: fa.EdgeKField[float],
 ):
-    _upwind_hflux_miura_cycl_stencil_03a(z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl, out=(p_out_e))
+    _average_horizontal_flux_subcycling_2(z_tracer_mflx_1_dsl, z_tracer_mflx_2_dsl, out=(p_out_e))