Add tracer advection to standalone driver [cycle 34]

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

@@ -5,12 +5,21 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
+from __future__ import annotations
 
 import dataclasses
+from typing import TYPE_CHECKING
 
 import gt4py.next as gtx
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
+from icon4py.model.common.utils import data_allocation as data_alloc
+
+
+if TYPE_CHECKING:
+    import gt4py.next.typing as gtx_typing
+
+    from icon4py.model.common.grid import icon as icon_grid
 
 
 @dataclasses.dataclass(frozen=True)
@@ -88,3 +97,26 @@ class AdvectionMetricState:
 
     #: vertical grid spacing at full levels
     ddqz_z_full: fa.CellKField[ta.wpfloat]
+
+
+def initialize_advection_diagnostic_state(
+    grid: icon_grid.IconGrid,
+    allocator: gtx_typing.FieldBufferAllocationUtil,
+) -> AdvectionDiagnosticState:
+    return AdvectionDiagnosticState(
+        airmass_now=data_alloc.zero_field(
+            grid, dims.CellDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat
+        ),
+        airmass_new=data_alloc.zero_field(
+            grid, dims.CellDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat
+        ),
+        grf_tend_tracer=data_alloc.zero_field(
+            grid, dims.CellDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat
+        ),
+        hfl_tracer=data_alloc.zero_field(
+            grid, dims.EdgeDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat
+        ),
+        vfl_tracer=data_alloc.zero_field(
+            grid, dims.CellDim, dims.KDim, allocator=allocator, dtype=ta.wpfloat
+        ),
+    )

model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py

@@ -317,6 +317,7 @@ def test_benchmark_solve_nonhydro(
         theta_v=data_alloc.random_field(mesh, dims.CellDim, dims.KDim, allocator=allocator),
         rho=data_alloc.random_field(mesh, dims.CellDim, dims.KDim, allocator=allocator),
         exner=data_alloc.random_field(mesh, dims.CellDim, dims.KDim, allocator=allocator),
+        tracer=[],
     )
     prognostic_state_nnew = prognostics.PrognosticState(
         w=data_alloc.random_field(
@@ -326,6 +327,7 @@ def test_benchmark_solve_nonhydro(
         theta_v=data_alloc.random_field(mesh, dims.CellDim, dims.KDim, allocator=allocator),
         rho=data_alloc.random_field(mesh, dims.CellDim, dims.KDim, allocator=allocator),
         exner=data_alloc.random_field(mesh, dims.CellDim, dims.KDim, allocator=allocator),
+        tracer=[],
     )
 
     prognostic_states = common_utils.TimeStepPair(prognostic_state_nnow, prognostic_state_nnew)

model/atmosphere/dycore/tests/dycore/utils.py

@@ -140,13 +140,15 @@ def create_prognostic_states(
         theta_v=sp.theta_v_now(),
         rho=sp.rho_now(),
         exner=sp.exner_now(),
+        tracer=[],
     )
     prognostic_state_nnew = prognostics.PrognosticState(
         w=sp.w_new(),
         vn=sp.vn_new(),
         theta_v=sp.theta_v_new(),
         rho=sp.rho_new(),
         exner=sp.exner_new(),
+        tracer=[],
     )
     prognostic_states = common_utils.TimeStepPair(prognostic_state_nnow, prognostic_state_nnew)
     return prognostic_states

model/common/src/icon4py/model/common/metrics/compute_advection_metrics.py

@@ -0,0 +1,53 @@
+# 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 types import ModuleType
+
+import numpy as np
+
+from icon4py.model.common import field_type_aliases as fa, type_alias as ta
+
+
+def compute_acvection_deepatmo_fields(
+    vct_a: fa.KField[ta.wpfloat],
+    nlev: int,
+    grid_sphere_radius: float,
+    array_ns: ModuleType = np,
+):
+    deepatmo_divh_mc = array_ns.zeros((nlev,))
+    deepatmo_divzU_mc = array_ns.zeros((nlev,))
+    deepatmo_divzL_mc = array_ns.zeros((nlev,))
+    height_uifc = vct_a[:nlev]
+    height_lifc = vct_a[1 : nlev + 1]
+    height_mc = 0.5 * (height_lifc + height_uifc)
+    radial_distance_mc = height_mc + grid_sphere_radius
+    radial_distance_lifc = grid_sphere_radius + height_lifc
+    radial_distance_uifc = grid_sphere_radius + height_uifc
+    deepatmo_gradh_mc = grid_sphere_radius / radial_distance_mc
+    deepatmo_divh_mc = (
+        deepatmo_gradh_mc
+        * 3.0
+        / 4.0
+        / (
+            1.0
+            - radial_distance_lifc
+            * radial_distance_uifc
+            / (radial_distance_lifc + radial_distance_uifc) ** 2
+        )
+    )
+    deepatmo_divzL_mc = 3.0 / (
+        1.0
+        + radial_distance_uifc / radial_distance_lifc
+        + (radial_distance_uifc / radial_distance_lifc) ** 2
+    )
+    deepatmo_divzU_mc = 3.0 / (
+        1.0
+        + radial_distance_lifc / radial_distance_uifc
+        + (radial_distance_lifc / radial_distance_uifc) ** 2
+    )
+    return deepatmo_divh_mc, deepatmo_divzL_mc, deepatmo_divzU_mc

model/common/src/icon4py/model/common/metrics/metrics_attributes.py

@@ -69,6 +69,9 @@
 ZD_INTCOEF_DSL: Final[str] = "zd_intcoef_dsl"
 ZD_VERTOFFSET_DSL: Final[str] = "zd_vertoffset_dsl"
 CELL_HEIGHT_ON_HALF_LEVEL: Final[str] = "vertical_coordinates_on_half_levels"
+DEEPATMO_DIVH: Final[str] = "deepatmo_divh_mc"
+DEEPATMO_DIVZL: Final[str] = "deepatmo_divzL_mc"
+DEEPATMO_DIVZU: Final[str] = "deepatmo_divzU_mc"
 
 
 attrs: dict[str, model.FieldMetaData] = {
@@ -463,4 +466,28 @@
         icon_var_name="z_ifc",
         dtype=ta.wpfloat,
     ),
+    DEEPATMO_DIVH: dict(
+        standard_name=DEEPATMO_DIVH,
+        long_name="",
+        units="",
+        dims=(dims.KDim),
+        icon_var_name="deepatmo_divh_mc",
+        dtype=ta.wpfloat,
+    ),
+    DEEPATMO_DIVZL: dict(
+        standard_name=DEEPATMO_DIVZL,
+        long_name="",
+        units="",
+        dims=(dims.KDim),
+        icon_var_name="deepatmo_divzL_mc",
+        dtype=ta.wpfloat,
+    ),
+    DEEPATMO_DIVZU: dict(
+        standard_name=DEEPATMO_DIVZU,
+        long_name="",
+        units="",
+        dims=(dims.KDim),
+        icon_var_name="deepatmo_divzU_mc",
+        dtype=ta.wpfloat,
+    ),
 }

model/common/src/icon4py/model/common/metrics/metrics_factory.py

@@ -33,6 +33,7 @@
 from icon4py.model.common.interpolation import interpolation_attributes, interpolation_factory
 from icon4py.model.common.interpolation.stencils import cell_2_edge_interpolation
 from icon4py.model.common.metrics import (
+    compute_advection_metrics,
     compute_coeff_gradekin,
     compute_diffusion_metrics,
     compute_zdiff_gradp_dsl,
@@ -952,6 +953,25 @@ def _register_computed_fields(self) -> None:  # noqa: PLR0915 [too-many-statemen
 
         self.register_provider(compute_diffusion_intcoef_and_vertoffset)
 
+        compute_advection_deepatmo_fields = factory.NumpyDataProvider(
+            func=functools.partial(
+                compute_advection_metrics.compute_advection_deepatmo_fields,
+                array_ns=self._xp,
+            ),
+            deps={
+                "vct_a": self._vertical_grid.interface_physical_height,
+            },
+            domain=(),
+            fields=(
+                attrs.DEEPATMO_DIVH,
+                attrs.DEEPATMO_DIVZL,
+                attrs.DEEPATMO_DIVZU,
+            ),
+            params={"nlev": self._grid.num_levels, "grid_sphere_radius": constants.EARTH_RADIUS},
+        )
+
+        self.register_provider(compute_advection_deepatmo_fields)
+
     @property
     def metadata(self) -> dict[str, model.FieldMetaData]:
         return self._attrs

model/common/src/icon4py/model/common/states/prognostic_state.py

@@ -35,6 +35,9 @@ class PrognosticState:
     ]  # horizontal wind normal to edges, vn(nproma, nlev, nblks_e)  [m/s]
     exner: fa.CellKField[ta.wpfloat]  # exner function, exner(nrpoma, nlev, nblks_c)
     theta_v: fa.CellKField[ta.wpfloat]  # virtual temperature, (nproma, nlev, nlbks_c) [K]
+    tracer: list[
+        fa.CellKField[ta.wpfloat]
+    ]  # tracer concentration (nproma,nlev,nblks_c,ntracer) [kg/kg]
 
     @property
     def w_1(self) -> fa.CellField[ta.wpfloat]:
@@ -44,6 +47,7 @@ def w_1(self) -> fa.CellField[ta.wpfloat]:
 def initialize_prognostic_state(
     grid: icon_grid.IconGrid,
     allocator: gtx_typing.FieldBufferAllocationUtil,
+    ntracer: int = 0,
 ) -> PrognosticState:
     """Initialize the prognostic state with zero fields."""
     rho = data_alloc.zero_field(
@@ -82,4 +86,14 @@ def initialize_prognostic_state(
         allocator=allocator,
         dtype=ta.wpfloat,
     )
-    return PrognosticState(rho=rho, w=w, vn=vn, exner=exner, theta_v=theta_v)
+    tracer = [
+        data_alloc.zero_field(
+            grid,
+            dims.CellDim,
+            dims.KDim,
+            allocator=allocator,
+            dtype=ta.wpfloat,
+        )
+        for _ in range(ntracer)
+    ]
+    return PrognosticState(rho=rho, w=w, vn=vn, exner=exner, theta_v=theta_v, tracer=tracer)

model/standalone_driver/src/icon4py/model/standalone_driver/config.py

@@ -34,3 +34,4 @@ class DriverConfig:
     vertical_cfl_threshold: ta.wpfloat = 0.85
     ndyn_substeps: int = 5
     enable_statistics_output: bool = False
+    ntracer: int = 0  # this is the default in ICON

model/standalone_driver/src/icon4py/model/standalone_driver/driver_states.py

@@ -17,6 +17,7 @@
 import devtools
 
 import icon4py.model.common.utils as common_utils
+from icon4py.model.atmosphere.advection import advection_states
 from icon4py.model.atmosphere.diffusion import diffusion_states
 from icon4py.model.atmosphere.dycore import dycore_states
 from icon4py.model.common import type_alias as ta
@@ -62,6 +63,7 @@ class DriverStates(NamedTuple):
 
     prep_advection_prognostic: dycore_states.PrepAdvection
     solve_nonhydro_diagnostic: dycore_states.DiagnosticStateNonHydro
+    tracer_advection_diagnostic: advection_states.AdvectionDiagnosticState
     diffusion_diagnostic: diffusion_states.DiffusionDiagnosticState
     prognostics: common_utils.TimeStepPair[prognostics.PrognosticState]
     diagnostic: diagnostics.DiagnosticState

model/standalone_driver/src/icon4py/model/standalone_driver/driver_utils.py

@@ -18,6 +18,7 @@
 import gt4py.next as gtx
 import gt4py.next.typing as gtx_typing
 
+from icon4py.model.atmosphere.advection import advection, advection_states
 from icon4py.model.atmosphere.diffusion import diffusion, diffusion_states
 from icon4py.model.atmosphere.dycore import dycore_states, solve_nonhydro as solve_nh
 from icon4py.model.common import (
@@ -156,14 +157,12 @@ def initialize_granules(
     vertical_grid: v_grid.VerticalGrid,
     diffusion_config: diffusion.DiffusionConfig,
     solve_nh_config: solve_nh.NonHydrostaticConfig,
+    advection_config: advection.AdvectionConfig,
     static_field_factories: driver_states.StaticFieldFactories,
     exchange: decomposition_defs.ExchangeRuntime,
     owner_mask: fa.CellField[bool],
     backend: model_backends.BackendLike,
-) -> tuple[
-    diffusion.Diffusion,
-    solve_nh.SolveNonhydro,
-]:
+) -> tuple[diffusion.Diffusion, solve_nh.SolveNonhydro, advection.Advection]:
     geometry_field_source = static_field_factories.geometry_field_source
     interpolation_field_source = static_field_factories.interpolation_field_source
     metrics_field_source = static_field_factories.metrics_field_source
@@ -342,7 +341,43 @@ def initialize_granules(
         owner_mask=owner_mask,
     )
 
-    return diffusion_granule, solve_nonhydro_granule
+    advection_granule = advection.convert_config_to_advection(
+        grid=grid,
+        backend=backend,
+        config=advection_config,
+        interpolation_state=advection_states.AdvectionInterpolationState(
+            geofac_div=interpolation_field_source.get(interpolation_attributes.GEOFAC_DIV),
+            rbf_vec_coeff_e=interpolation_field_source.get(
+                interpolation_attributes.RBF_VEC_COEFF_E
+            ),
+            pos_on_tplane_e_1=interpolation_field_source.get(
+                interpolation_attributes.POS_ON_TPLANE_E_X
+            ),
+            pos_on_tplane_e_2=interpolation_field_source.get(
+                interpolation_attributes.POS_ON_TPLANE_E_Y
+            ),
+        ),
+        least_squares_state=advection_states.AdvectionLeastSquaresState(
+            # TODO(ricoh): [c34] integrate #1028 after that's merged
+            lsq_pseudoinv_1=data_alloc.constant_field(
+                grid, 0.0, dims.CellDim, dims.C2E2CDim, allocator=backend
+            ),
+            lsq_pseudoinv_2=data_alloc.constant_field(
+                grid, 0.0, dims.CellDim, dims.C2E2CDim, allocator=backend
+            ),
+        ),
+        metric_state=advection_states.AdvectionMetricState(
+            deepatmo_divh=metrics_field_source.get(metrics_attributes.DEEPATMO_DIVH),
+            deepatmo_divzl=metrics_field_source.get(metrics_attributes.DEEPATMO_DIVZL),
+            deepatmo_divzu=metrics_field_source.get(metrics_attributes.DEEPATMO_DIVZU),
+            ddqz_z_full=metrics_field_source.get(metrics_attributes.DDQZ_Z_FULL),
+        ),
+        edge_params=edge_geometry,
+        cell_params=cell_geometry,
+        exchange=exchange,
+    )
+
+    return diffusion_granule, solve_nonhydro_granule, advection_granule
 
 
 def find_maximum_from_field(
@@ -430,6 +465,7 @@ def display_driver_setup_in_log_file(
     log.info(f"Vertical CFL threshold : {config.vertical_cfl_threshold}")
     log.info(f"Second-order divdamp   : {config.apply_extra_second_order_divdamp}")
     log.info(f"Statistics enabled     : {config.enable_statistics_output}")
+    log.info(f"Number of tracers      : {config.ntracer}")
     log.info("")
 
     log.info("==== Vertical Grid Parameters ====")