Add tracer advection to standalone driver [cycle 34]

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,44 @@ 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] initialize with zero, check if works
+            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(
+            # TODO(ricoh): [c34] DEEPATMO_XYZ don't exist yet
+            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(

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

@@ -16,6 +16,7 @@
 from gt4py.next import config as gtx_config, metrics as gtx_metrics
 
 import icon4py.model.common.utils as common_utils
+from icon4py.model.atmosphere.advection import advection
 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 dimension as dims, model_backends, model_options, type_alias as ta
@@ -46,6 +47,7 @@ def __init__(
         static_field_factories: driver_states.StaticFieldFactories,
         diffusion_granule: diffusion.Diffusion,
         solve_nonhydro_granule: solve_nh.SolveNonhydro,
+        advection_granule: advection.Advection,  # NOTE(ricoh): [c34] wip
     ):
         self.config = config
         self.backend = backend
@@ -54,6 +56,7 @@ def __init__(
         self.diffusion = diffusion_granule
         self.solve_nonhydro = solve_nonhydro_granule
         self.model_time_variables = driver_states.ModelTimeVariables(config=config)
+        self.tracer_advection = advection_granule  # NOTE (ricoh): [c34] wip
         self.timer_collection = driver_states.TimerCollection(
             [timer.value for timer in driver_states.DriverTimers]
         )
@@ -175,6 +178,9 @@ def _integrate_one_time_step(
                 )
             timer_diffusion.capture()
 
+        # NOTE (ricoh): [c34] something like
+        # NOTE (ricoh): self.tracer_advection.run(...) # pass list of tracers (or by default step all of the initialized ones)
+
         prognostic_states.swap()
 
     def _update_time_levels_for_velocity_tendencies(
@@ -470,6 +476,7 @@ def _read_config(
     driver_config.DriverConfig,
     v_grid.VerticalGridConfig,
     diffusion.DiffusionConfig,
+    advection.AdvectionConfig,
     solve_nh.NonHydrostaticConfig,
 ]:
     vertical_grid_config = v_grid.VerticalGridConfig(
@@ -492,6 +499,14 @@ def _read_config(
         velocity_boundary_diffusion_denom=200.0,
     )
 
+    # TODO (ricoh): [c34] check config correctness
+    advection_config = advection.AdvectionConfig(
+        horizontal_advection_limiter=advection.HorizontalAdvectionLimiter.NO_LIMITER,
+        horizontal_advection_type=advection.HorizontalAdvectionType.LINEAR_2ND_ORDER,
+        vertical_advection_limiter=advection.VerticalAdvectionLimiter.NO_LIMITER,
+        vertical_advection_type=advection.VerticalAdvectionType.UPWIND_1ST_ORDER,
+    )
+
     nonhydro_config = solve_nh.NonHydrostaticConfig(
         fourth_order_divdamp_factor=0.0025,
     )
@@ -514,6 +529,7 @@ def _read_config(
         icon4py_driver_config,
         vertical_grid_config,
         diffusion_config,
+        advection_config,
         nonhydro_config,
     )
 
@@ -569,9 +585,11 @@ def initialize_driver(
     allocator = model_backends.get_allocator(backend)
 
     log.info("Initializing the driver")
-    driver_config, vertical_grid_config, diffusion_config, solve_nh_config = _read_config(
-        output_path=output_path,
-        enable_profiling=False,
+    driver_config, vertical_grid_config, diffusion_config, advection_config, solve_nh_config = (
+        _read_config(
+            output_path=output_path,
+            enable_profiling=False,
+        )
     )
 
     log.info(f"initializing the grid manager from '{grid_file_path}'")
@@ -616,11 +634,13 @@ def initialize_driver(
     (
         diffusion_granule,
         solve_nonhydro_granule,
+        advection_granule,  # NOTE(ricoh): [c34] wip
     ) = driver_utils.initialize_granules(
         grid=grid_manager.grid,
         vertical_grid=vertical_grid,
         diffusion_config=diffusion_config,
         solve_nh_config=solve_nh_config,
+        advection_config=advection_config,
         static_field_factories=static_field_factories,
         exchange=exchange,
         owner_mask=gtx.as_field(
@@ -637,6 +657,7 @@ def initialize_driver(
         static_field_factories=static_field_factories,
         diffusion_granule=diffusion_granule,
         solve_nonhydro_granule=solve_nonhydro_granule,
+        advection_granule=advection_granule,
     )
 
     return icon4py_driver