Enable optional stencils and modules in Liskov

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

@@ -27,14 +27,14 @@ def _hor_adv_stencil_01(
     geofac_div: Field[[CEDim], float],
     p_dtime: float,
 ) -> Field[[CellDim, KDim], float]:
-    tracer_new = (
+    tracer_new_hor = (
         tracer_now * rhodz_now
         - p_dtime
         * deepatmo_divh
         * neighbor_sum(p_mflx_tracer_h(C2E) * geofac_div(C2CE), axis=C2EDim)
     ) / rhodz_new
 
-    return tracer_new
+    return tracer_new_hor
 
 
 @program
@@ -45,7 +45,7 @@ def hor_adv_stencil_01(
     rhodz_now: Field[[CellDim, KDim], float],
     rhodz_new: Field[[CellDim, KDim], float],
     geofac_div: Field[[CEDim], float],
-    tracer_new: Field[[CellDim, KDim], float],
+    tracer_new_hor: Field[[CellDim, KDim], float],
     p_dtime: float,
 ):
     _hor_adv_stencil_01(
@@ -56,5 +56,5 @@ def hor_adv_stencil_01(
         rhodz_new,
         geofac_div,
         p_dtime,
-        out=tracer_new,
+        out=tracer_new_hor,
     )

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

@@ -33,6 +33,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl(
     dual_normal_cell_2: Field[[ECDim], wpfloat],
     p_dthalf: wpfloat,
 ) -> tuple[
+    Field[[EdgeDim, KDim], int32],
     Field[[EdgeDim, KDim], int32],
     Field[[EdgeDim, KDim], int32],
     Field[[EdgeDim, KDim], vpfloat],
@@ -41,6 +42,7 @@ def _mo_advection_traj_btraj_compute_o1_dsl(
     lvn_pos = where(p_vn > 0.0, True, False)
 
     p_cell_idx = where(lvn_pos, cell_idx(E2EC[0]), cell_idx(E2EC[1]))
+    p_cell_rel_idx_dsl = where(lvn_pos, int32(0), int32(1))
     p_cell_blk = where(lvn_pos, cell_blk(E2EC[0]), cell_blk(E2EC[1]))
 
     z_ntdistv_bary_1 = -(
@@ -67,7 +69,13 @@ def _mo_advection_traj_btraj_compute_o1_dsl(
         + z_ntdistv_bary_2 * dual_normal_cell_2(E2EC[1]),
     )
 
-    return p_cell_idx, p_cell_blk, astype(p_distv_bary_1, vpfloat), astype(p_distv_bary_2, vpfloat)
+    return (
+        p_cell_idx,
+        p_cell_rel_idx_dsl,
+        p_cell_blk,
+        astype(p_distv_bary_1, vpfloat),
+        astype(p_distv_bary_2, vpfloat),
+    )
 
 
 @program(grid_type=GridType.UNSTRUCTURED)
@@ -83,6 +91,7 @@ def mo_advection_traj_btraj_compute_o1_dsl(
     primal_normal_cell_2: Field[[ECDim], wpfloat],
     dual_normal_cell_2: Field[[ECDim], wpfloat],
     p_cell_idx: Field[[EdgeDim, KDim], int32],
+    p_cell_rel_idx_dsl: Field[[EdgeDim, KDim], int32],
     p_cell_blk: Field[[EdgeDim, KDim], int32],
     p_distv_bary_1: Field[[EdgeDim, KDim], vpfloat],
     p_distv_bary_2: Field[[EdgeDim, KDim], vpfloat],
@@ -104,7 +113,13 @@ def mo_advection_traj_btraj_compute_o1_dsl(
         primal_normal_cell_2,
         dual_normal_cell_2,
         p_dthalf,
-        out=(p_cell_idx, p_cell_blk, p_distv_bary_1, p_distv_bary_2),
+        out=(
+            p_cell_idx,
+            p_cell_rel_idx_dsl,
+            p_cell_blk,
+            p_distv_bary_1,
+            p_distv_bary_2,
+        ),
         domain={
             EdgeDim: (horizontal_start, horizontal_end),
             KDim: (vertical_start, vertical_end),

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

@@ -24,5 +24,5 @@ def _set_zero_c_k() -> Field[[CellDim, KDim], float]:
 
 
 @program(grid_type=GridType.UNSTRUCTURED)
-def set_zero_c_k(field: Field[[CellDim, KDim], float]):
-    _set_zero_c_k(out=field)
+def set_zero_c_k(z_cfl: Field[[CellDim, KDim], float]):
+    _set_zero_c_k(out=z_cfl)

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

@@ -12,13 +12,13 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 from gt4py.next import GridType
 from gt4py.next.ffront.decorator import field_operator, program
-from gt4py.next.ffront.fbuiltins import Field
+from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where
 
 from icon4py.model.common.dimension import CellDim, KDim, Koff
 
 
 @field_operator
-def _vert_adv_stencil_01(
+def _vert_adv_stencil_01a(
     tracer_now: Field[[CellDim, KDim], float],
     rhodz_now: Field[[CellDim, KDim], float],
     p_mflx_tracer_v: Field[[CellDim, KDim], float],
@@ -35,6 +35,42 @@ def _vert_adv_stencil_01(
     return tracer_new
 
 
+@field_operator
+def _vert_adv_stencil_01(
+    tracer_now: Field[[CellDim, KDim], float],
+    rhodz_now: Field[[CellDim, KDim], float],
+    p_mflx_tracer_v: Field[[CellDim, KDim], float],
+    deepatmo_divzl: Field[[KDim], float],
+    deepatmo_divzu: Field[[KDim], float],
+    rhodz_new: Field[[CellDim, KDim], float],
+    k: Field[[KDim], int32],
+    p_dtime: float,
+    ivadv_tracer: int32,
+    iadv_slev_jt: int32,
+) -> Field[[CellDim, KDim], float]:
+    k = broadcast(k, (CellDim, KDim))
+
+    tracer_new = (
+        where(
+            (iadv_slev_jt <= k),
+            _vert_adv_stencil_01a(
+                tracer_now,
+                rhodz_now,
+                p_mflx_tracer_v,
+                deepatmo_divzl,
+                deepatmo_divzu,
+                rhodz_new,
+                p_dtime,
+            ),
+            tracer_now,
+        )
+        if (ivadv_tracer != int32(0))
+        else tracer_now
+    )
+
+    return tracer_new
+
+
 @program(grid_type=GridType.UNSTRUCTURED)
 def vert_adv_stencil_01(
     tracer_now: Field[[CellDim, KDim], float],
@@ -43,8 +79,11 @@ def vert_adv_stencil_01(
     deepatmo_divzl: Field[[KDim], float],
     deepatmo_divzu: Field[[KDim], float],
     rhodz_new: Field[[CellDim, KDim], float],
-    tracer_new: Field[[CellDim, KDim], float],
+    k: Field[[KDim], int32],
     p_dtime: float,
+    ivadv_tracer: int32,
+    iadv_slev_jt: int32,
+    tracer_new: Field[[CellDim, KDim], float],
 ):
     _vert_adv_stencil_01(
         tracer_now,
@@ -53,6 +92,9 @@ def vert_adv_stencil_01(
         deepatmo_divzl,
         deepatmo_divzu,
         rhodz_new,
+        k,
         p_dtime,
+        ivadv_tracer,
+        iadv_slev_jt,
         out=tracer_new,
     )

model/atmosphere/advection/tests/stencil_tests/test_vert_adv_stencil_01.py

@@ -13,10 +13,12 @@
 
 import numpy as np
 import pytest
+from gt4py.next import as_field
+from gt4py.next.ffront.fbuiltins import int32
 
 from icon4py.model.atmosphere.advection.vert_adv_stencil_01 import vert_adv_stencil_01
 from icon4py.model.common.dimension import CellDim, KDim
-from icon4py.model.common.test_utils.helpers import StencilTest, random_field, zero_field
+from icon4py.model.common.test_utils.helpers import StencilTest, _shape, random_field, zero_field
 
 
 class TestVertAdvStencil01(StencilTest):
@@ -32,14 +34,29 @@ def reference(
         deepatmo_divzl: np.array,
         deepatmo_divzu: np.array,
         rhodz_new: np.array,
-        p_dtime,
+        k: np.array,
+        ivadv_tracer: int32,
+        iadv_slev_jt: int32,
+        p_dtime: np.float64,
         **kwargs,
     ) -> np.array:
-        tracer_new = (
-            tracer_now * rhodz_now
-            + p_dtime
-            * (p_mflx_tracer_v[:, 1:] * deepatmo_divzl - p_mflx_tracer_v[:, :-1] * deepatmo_divzu)
-        ) / rhodz_new
+        if ivadv_tracer != int32(0):
+            tracer_new = np.where(
+                (iadv_slev_jt <= k),
+                (
+                    tracer_now * rhodz_now
+                    + p_dtime
+                    * (
+                        p_mflx_tracer_v[:, 1:] * deepatmo_divzl
+                        - p_mflx_tracer_v[:, :-1] * deepatmo_divzu
+                    )
+                )
+                / rhodz_new,
+                tracer_now,
+            )
+        else:
+            tracer_new = tracer_now
+
         return dict(tracer_new=tracer_new)
 
     @pytest.fixture
@@ -50,7 +67,10 @@ def input_data(self, grid):
         deepatmo_divzl = random_field(grid, KDim)
         deepatmo_divzu = random_field(grid, KDim)
         rhodz_new = random_field(grid, CellDim, KDim)
+        k = as_field((KDim,), np.arange(0, _shape(grid, KDim)[0], dtype=int32))
         p_dtime = np.float64(5.0)
+        ivadv_tracer = int32(1)
+        iadv_slev_jt = int32(4)
         tracer_new = zero_field(grid, CellDim, KDim)
         return dict(
             tracer_now=tracer_now,
@@ -59,6 +79,9 @@ def input_data(self, grid):
             deepatmo_divzl=deepatmo_divzl,
             deepatmo_divzu=deepatmo_divzu,
             rhodz_new=rhodz_new,
+            k=k,
             p_dtime=p_dtime,
+            ivadv_tracer=ivadv_tracer,
+            iadv_slev_jt=iadv_slev_jt,
             tracer_new=tracer_new,
         )

model/tox.ini

@@ -30,11 +30,13 @@ allowlist_externals =
 commands =
     pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-skip -n auto {posargs}
     pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/dycore/tests/dycore_stencil_tests --benchmark-skip -n auto {posargs}
+    pytest -v -s -m "not slow_tests" --cov --cov-append atmosphere/advection/tests/stencil_tests --benchmark-skip -n auto {posargs}
 
 [testenv:run_benchmarks]
 commands =
     pytest -s -m "not slow_tests" atmosphere/diffusion/tests/diffusion_stencil_tests --benchmark-only {posargs}
     pytest -s -m "not slow_tests" atmosphere/dycore/tests/dycore_stencil_tests --benchmark-only {posargs}
+    pytest -s -m "not slow_tests" atmosphere/advection/tests/stencil_tests --benchmark-only {posargs}
 
 
 [testenv:dev]

tools/README.md

@@ -148,6 +148,8 @@ In addition, other optional keyword arguments are the following:
 
 - `copies`: Takes a boolean string input, and controls whether before field copies should be made or not. If set to False only the `#ifdef __DSL_VERIFY` directive is generated. Defaults to true.<br><br>
 
+- `optional_module`: Takes a boolean string input, and controls whether stencils is part of an optional module. Defaults to "None".<br><br>
+
 #### `!$DSL END STENCIL()`
 
 This directive denotes the end of a stencil. The required argument is `name`, which must match the name of the preceding `START STENCIL` directive.

tools/src/icon4pytools/liskov/cli.py

@@ -56,6 +56,12 @@ def main(ctx: click.Context) -> None:
     default=True,
     help="Adds fused or unfused stencils.",
 )
+@click.option(
+    "--enable-dsl-optional",
+    type=click.Choice(["advection", "graupel", "no"]),
+    default="no",
+    help="Specify the DSL optional module to enable.",
+)
 @click.argument(
     "input_path",
     type=click.Path(exists=True, dir_okay=False, resolve_path=True, path_type=pathlib.Path),
@@ -70,10 +76,11 @@ def integrate(
     fused: bool,
     profile: bool,
     metadatagen: bool,
+    enable_dsl_optional: str,
 ) -> None:
     mode = "integration"
     iface = parse_fortran_file(input_path, output_path, mode)
-    iface_gt4py = process_stencils(iface, fused)
+    iface_gt4py = process_stencils(iface, fused, optional_modules_to_enable=enable_dsl_optional)
     run_code_generation(
         input_path,
         output_path,

tools/src/icon4pytools/liskov/codegen/integration/deserialise.py

@@ -291,7 +291,6 @@ def create_stencil_data(
             bounds = self._make_bounds(named_args)
             fields = self._make_fields(named_args, field_dimensions)
             fields_w_tolerance = self._update_tolerances(named_args, fields)
-
             deserialised.append(
                 dtype(
                     name=stencil_name,
@@ -305,14 +304,21 @@ def create_stencil_data(
         return deserialised
 
     def _pop_additional_attributes(
-        self, dtype: Type[StartStencilData | StartFusedStencilData], named_args: dict[str, Any]
+        self,
+        dtype: Type[StartStencilData | StartFusedStencilData],
+        named_args: dict[str, Any],
     ) -> dict:
         """Pop and return additional attributes specific to StartStencilData."""
         additional_attrs = {}
         if dtype == StartStencilData:
             mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false"))
             copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true"))
-            additional_attrs = {"mergecopy": mergecopy, "copies": copies}
+            optional_module = pop_item_from_dict(named_args, "optional_module", "None")
+            additional_attrs = {
+                "mergecopy": mergecopy,
+                "copies": copies,
+                "optional_module": optional_module,
+            }
         return additional_attrs
 
     @staticmethod

tools/src/icon4pytools/liskov/codegen/integration/generate.py

@@ -143,6 +143,7 @@ def _generate_start_stencil(self) -> None:
                     acc_present=stencil.acc_present,
                     mergecopy=stencil.mergecopy,
                     copies=stencil.copies,
+                    optional_module=stencil.optional_module,
                 )
                 i += 2
 

tools/src/icon4pytools/liskov/codegen/integration/interface.py

@@ -89,6 +89,7 @@ class BaseStartStencilData(CodeGenInput):
 class StartStencilData(BaseStartStencilData):
     mergecopy: Optional[bool]
     copies: Optional[bool]
+    optional_module: Optional[str]
 
 
 @dataclass