CI for standalone_driver

.github/workflows/icon4py-test-model.yml

@@ -23,7 +23,7 @@ jobs:
       matrix:
         python-version: ["3.10", "3.11"]
         backend: ["embedded", "dace_cpu", "gtfn_cpu"]
-        component: ["advection", "diffusion", "dycore", "microphysics", "muphys", "driver", "common"]
+        component: ["advection", "diffusion", "dycore", "microphysics", "muphys", "driver", "common", "standalone_driver"]
     steps:
       - name: Checkout
         uses: actions/checkout@v6

ci/default.yml

@@ -52,6 +52,9 @@ test_tools_datatests_aarch64:
     - if: $COMPONENT == 'driver' || $COMPONENT == 'dycore' || $COMPONENT == 'muphys'
       variables:
         SLURM_TIMELIMIT: '00:30:00'
+    - if: $COMPONENT == 'standalone_driver'
+      variables:
+        SLURM_TIMELIMIT: '04:00:00' # TODO: reset to 45 minutes once time limit estimation for dace_gpu is complete
     - when: on_success
       variables:
         SLURM_TIMELIMIT: '00:15:00'

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/apply_rayleigh_damping_mechanism.py

@@ -15,19 +15,17 @@
 @gtx.field_operator
 def _apply_rayleigh_damping_mechanism(
     z_raylfac: fa.KField[wpfloat],
-    w_1: fa.CellField[wpfloat],
     w: fa.CellKField[wpfloat],
 ) -> fa.CellKField[wpfloat]:
     """Formerly known as _mo_solve_nonhydro_stencil_54."""
     z_raylfac = broadcast(z_raylfac, (dims.CellDim, dims.KDim))
-    w_wp = z_raylfac * w + (wpfloat("1.0") - z_raylfac) * w_1
+    w_wp = z_raylfac * w
     return w_wp
 
 
 @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
 def apply_rayleigh_damping_mechanism(
     z_raylfac: fa.KField[wpfloat],
-    w_1: fa.CellField[wpfloat],
     w: fa.CellKField[wpfloat],
     horizontal_start: gtx.int32,
     horizontal_end: gtx.int32,
@@ -36,7 +34,6 @@ def apply_rayleigh_damping_mechanism(
 ) -> None:
     _apply_rayleigh_damping_mechanism(
         z_raylfac,
-        w_1,
         w,
         out=w,
         domain={

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/stencils/vertically_implicit_dycore_solver.py

@@ -343,13 +343,11 @@ def _vertically_implicit_solver_at_predictor_step(
         cpd=dycore_consts.cpd,
     )
 
-    w_1 = broadcast(wpfloat("0.0"), (dims.CellDim,))
     if rayleigh_type == rayleigh_damping_options.KLEMP:
         next_w = concat_where(
             (dims.KDim > 0) & (dims.KDim < end_index_of_damping_layer + 1),
             _apply_rayleigh_damping_mechanism(
                 z_raylfac=rayleigh_damping_factor,
-                w_1=w_1,
                 w=next_w,
             ),
             next_w,
@@ -658,14 +656,11 @@ def _vertically_implicit_solver_at_corrector_step(
         cpd=dycore_consts.cpd,
     )
 
-    w_1 = broadcast(wpfloat("0.0"), (dims.CellDim,))
-
     if rayleigh_type == rayleigh_damping_options.KLEMP:
         next_w = concat_where(
             (dims.KDim > 0) & (dims.KDim < end_index_of_damping_layer + 1),
             _apply_rayleigh_damping_mechanism(
                 z_raylfac=rayleigh_damping_factor,
-                w_1=w_1,
                 w=next_w,
             ),
             next_w,

model/atmosphere/dycore/tests/dycore/stencil_tests/test_apply_rayleigh_damping_mechanism.py

@@ -25,12 +25,10 @@
 def apply_rayleigh_damping_mechanism_numpy(
     connectivities: dict[gtx.Dimension, np.ndarray],
     z_raylfac: np.ndarray,
-    w_1: np.ndarray,
     w: np.ndarray,
 ) -> np.ndarray:
     z_raylfac = np.expand_dims(z_raylfac, axis=0)
-    w_1 = np.expand_dims(w_1, axis=-1)
-    w = z_raylfac * w + (1.0 - z_raylfac) * w_1
+    w = z_raylfac * w
     return w
 
 
@@ -42,22 +40,19 @@ class TestApplyRayleighDampingMechanism(StencilTest):
     def reference(
         connectivities: dict[gtx.Dimension, np.ndarray],
         z_raylfac: np.ndarray,
-        w_1: np.ndarray,
         w: np.ndarray,
         **kwargs: Any,
     ) -> dict:
-        w = apply_rayleigh_damping_mechanism_numpy(connectivities, z_raylfac, w_1, w)
+        w = apply_rayleigh_damping_mechanism_numpy(connectivities, z_raylfac, w)
         return dict(w=w)
 
     @pytest.fixture
     def input_data(self, grid: base.Grid) -> dict[str, gtx.Field | state_utils.ScalarType]:
         z_raylfac = random_field(grid, dims.KDim, dtype=wpfloat)
         w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat)
-        w_1 = w[dims.KDim(0)]
 
         return dict(
             z_raylfac=z_raylfac,
-            w_1=w_1,
             w=w,
             horizontal_start=0,
             horizontal_end=gtx.int32(grid.num_cells),

model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_corrector_step.py

@@ -299,7 +299,6 @@ def reference(
             next_w[:, :n_lev],
         )
 
-        w_1 = next_w[:, 0]
         if rayleigh_type == constants.RayleighType.KLEMP:
             next_w[:, :n_lev] = np.where(
                 (horizontal_start <= horz_idx)
@@ -309,7 +308,6 @@ def reference(
                 apply_rayleigh_damping_mechanism_numpy(
                     connectivities=connectivities,
                     z_raylfac=rayleigh_damping_factor,
-                    w_1=w_1,
                     w=next_w[:, :n_lev],
                 ),
                 next_w[:, :n_lev],

model/atmosphere/dycore/tests/dycore/stencil_tests/test_vertically_implicit_dycore_solver_at_predictor_step.py

@@ -328,7 +328,6 @@ def reference(
             next_w[:, :n_lev],
         )
 
-        w_1 = next_w[:, 0]
         if rayleigh_type == constants.RayleighType.KLEMP:
             next_w[:, :n_lev] = np.where(
                 (start_cell_index_nudging <= horz_idx)
@@ -338,7 +337,6 @@ def reference(
                 apply_rayleigh_damping_mechanism_numpy(
                     connectivities=connectivities,
                     z_raylfac=rayleigh_damping_factor,
-                    w_1=w_1,
                     w=next_w[:, :n_lev],
                 ),
                 next_w[:, :n_lev],

model/common/src/icon4py/model/common/grid/vertical.py

@@ -529,7 +529,8 @@ def _compute_vct_a_and_vct_b(  # noqa: PLR0912 [too-many-branches]
 
 
 def get_vct_a_and_vct_b(
-    vertical_config: VerticalGridConfig, allocator: gtx_typing.Allocator
+    vertical_config: VerticalGridConfig,
+    allocator: gtx_typing.Allocator,
 ) -> tuple[fa.KField, fa.KField]:
     """
     get vct_a and vct_b.

model/common/src/icon4py/model/common/interpolation/interpolation_fields.py

@@ -13,7 +13,6 @@
 from typing import Final
 
 import numpy as np
-import scipy
 from gt4py import next as gtx
 from gt4py.next import where
 
@@ -1163,11 +1162,12 @@ def compute_lsq_pseudoinv(
     min_rlcell_int: int,
     lsq_dim_unk: int,
     lsq_dim_c: int,
+    array_ns: ModuleType = np,
 ) -> data_alloc.NDArray:
     for jjb in range(lsq_dim_c):
         for jjk in range(lsq_dim_unk):
             for jc in range(start_idx, min_rlcell_int):
-                u, s, v_t, _ = scipy.linalg.lapack.dgesdd(z_lsq_mat_c[jc, :, :])
+                u, s, v_t = array_ns.linalg.svd(z_lsq_mat_c[jc, :, :])
                 if cell_owner_mask[jc]:
                     lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = (
                         lsq_pseudoinv[jc, :lsq_dim_unk, jjb]
@@ -1181,11 +1181,12 @@ def compute_lsq_weights_c(
     lsq_weights_c_jc: data_alloc.NDArray,
     lsq_dim_stencil: int,
     lsq_wgt_exp: int,
+    array_ns: ModuleType = np,
 ) -> data_alloc.NDArray:
     for js in range(lsq_dim_stencil):
-        z_norm = np.sqrt(np.dot(z_dist_g[js, :], z_dist_g[js, :]))
+        z_norm = array_ns.sqrt(array_ns.dot(z_dist_g[js, :], z_dist_g[js, :]))
         lsq_weights_c_jc[js] = 1.0 / (z_norm**lsq_wgt_exp)
-    return lsq_weights_c_jc / np.max(lsq_weights_c_jc)
+    return lsq_weights_c_jc / array_ns.max(lsq_weights_c_jc)
 
 
 def compute_z_lsq_mat_c(
@@ -1234,18 +1235,18 @@ def compute_lsq_coeffs(
     match base_grid.GeometryType(geometry_type):
         case base_grid.GeometryType.ICOSAHEDRON:
             for js in range(lsq_dim_stencil):
-                z_dist_g[:, js, :] = np.asarray(
+                z_dist_g[:, js, :] = array_ns.asarray(
                     gnomonic_proj(
-                        cell_lon[:], cell_lat[:], cell_lon[c2e2c[:, js]], cell_lat[c2e2c[:, js]]
+                        cell_lon, cell_lat, cell_lon[c2e2c[:, js]], cell_lat[c2e2c[:, js]]
                     )
                 ).T
 
             z_dist_g *= grid_sphere_radius
             min_lsq_bound = min(lsq_dim_unk, lsq_dim_c)
-
             for jc in range(start_idx, min_rlcell_int):
                 if cell_owner_mask[jc]:
                     z_lsq_mat_c[jc, :min_lsq_bound, :min_lsq_bound] = 1.0
+
         case base_grid.GeometryType.TORUS:
             for jc in range(start_idx, min_rlcell_int):
                 ilc_s = c2e2c[jc, :lsq_dim_stencil]
@@ -1265,7 +1266,7 @@ def compute_lsq_coeffs(
 
     for jc in range(start_idx, min_rlcell_int):
         lsq_weights_c[jc, :] = compute_lsq_weights_c(
-            z_dist_g[jc, :, :], lsq_weights_c[jc, :], lsq_dim_stencil, lsq_wgt_exp
+            z_dist_g[jc, :, :], lsq_weights_c[jc, :], lsq_dim_stencil, lsq_wgt_exp, array_ns
         )
         z_lsq_mat_c[jc, js, :lsq_dim_unk] = compute_z_lsq_mat_c(
             cell_owner_mask,
@@ -1288,6 +1289,7 @@ def compute_lsq_coeffs(
         min_rlcell_int,
         lsq_dim_unk,
         lsq_dim_c,
+        array_ns,
     )
     exchange(lsq_pseudoinv[:, 0, :])
     exchange(lsq_pseudoinv[:, 1, :])

model/common/src/icon4py/model/common/interpolation/stencils/edge_2_cell_vector_rbf_interpolation.py

@@ -9,14 +9,14 @@
 from gt4py.next import neighbor_sum
 
 from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
-from icon4py.model.common.dimension import C2E2C2E, C2E2C2EDim
+from icon4py.model.common.dimension import C2E2C2E
 
 
 @gtx.field_operator
 def _edge_2_cell_vector_rbf_interpolation(
     p_e_in: fa.EdgeKField[ta.wpfloat],
-    ptr_coeff_1: gtx.Field[gtx.Dims[dims.CellDim, C2E2C2EDim], ta.wpfloat],
-    ptr_coeff_2: gtx.Field[gtx.Dims[dims.CellDim, C2E2C2EDim], ta.wpfloat],
+    ptr_coeff_1: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2C2EDim], ta.wpfloat],
+    ptr_coeff_2: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2C2EDim], ta.wpfloat],
 ) -> tuple[fa.CellKField[ta.wpfloat], fa.CellKField[ta.wpfloat]]:
     """
     Performs vector RBF reconstruction at cell center from edge center.
@@ -33,16 +33,16 @@ def _edge_2_cell_vector_rbf_interpolation(
     Returns:
         RBF reconstructed vector at cell center.
     """
-    p_u_out = neighbor_sum(ptr_coeff_1 * p_e_in(C2E2C2E), axis=C2E2C2EDim)
-    p_v_out = neighbor_sum(ptr_coeff_2 * p_e_in(C2E2C2E), axis=C2E2C2EDim)
+    p_u_out = neighbor_sum(ptr_coeff_1 * p_e_in(C2E2C2E), axis=dims.C2E2C2EDim)
+    p_v_out = neighbor_sum(ptr_coeff_2 * p_e_in(C2E2C2E), axis=dims.C2E2C2EDim)
     return p_u_out, p_v_out
 
 
 @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
 def edge_2_cell_vector_rbf_interpolation(
     p_e_in: fa.EdgeKField[ta.wpfloat],
-    ptr_coeff_1: gtx.Field[gtx.Dims[dims.CellDim, C2E2C2EDim], ta.wpfloat],
-    ptr_coeff_2: gtx.Field[gtx.Dims[dims.CellDim, C2E2C2EDim], ta.wpfloat],
+    ptr_coeff_1: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2C2EDim], ta.wpfloat],
+    ptr_coeff_2: gtx.Field[gtx.Dims[dims.CellDim, dims.C2E2C2EDim], ta.wpfloat],
     p_u_out: fa.CellKField[ta.wpfloat],
     p_v_out: fa.CellKField[ta.wpfloat],
     horizontal_start: gtx.int32,

model/driver/src/icon4py/model/driver/icon4py_configuration.py

@@ -115,7 +115,7 @@ def _jabw_diffusion_config(n_substeps: int):
             hdiff_efdt_ratio=10.0,
             hdiff_w_efdt_ratio=15.0,
             smagorinski_scaling_factor=0.025,
-            zdiffu_t=True,
+            zdiffu_t=False,
             velocity_boundary_diffusion_denom=200.0,
         )
 
@@ -124,6 +124,7 @@ def _jabw_nonhydro_config():
             # original igradp_method is 2
             # original divdamp_order is 4
             fourth_order_divdamp_factor=0.0025,
+            rayleigh_coeff=0.1,
         )
 
     def _mch_ch_r04b09_config():
@@ -144,7 +145,7 @@ def _mch_ch_r04b09_config():
     def _jablonowski_williamson_config():
         icon_run_config = Icon4pyRunConfig(
             dtime=datetime.timedelta(seconds=300.0),
-            end_date=datetime.datetime(1, 1, 1, 0, 30, 0),
+            end_date=datetime.datetime(1, 1, 1, 0, 5, 0),
             apply_initial_stabilization=False,
             n_substeps=5,
             backend=backend,

model/driver/tests/driver/integration_tests/test_icon4py.py

@@ -82,6 +82,19 @@
             False,
             False,
         ),
+        (
+            definitions.Experiments.JW,
+            1,
+            2,
+            1,
+            5,
+            "2008-09-01T00:00:00.000",
+            "2008-09-01T00:05:00.000",
+            "2008-09-01T00:05:00.000",
+            "2008-09-01T00:05:00.000",
+            False,
+            False,
+        ),
     ],
 )
 def test_run_timeloop_single_step(
@@ -106,9 +119,14 @@ def test_run_timeloop_single_step(
     savepoint_nonhydro_exit: sb.IconNonHydroExitSavepoint,
     backend: gtx_typing.Backend,
 ):
-    if experiment == definitions.Experiments.GAUSS3D:
+    if experiment in (definitions.Experiments.GAUSS3D, definitions.Experiments.JW):
+        experiment_type = (
+            driver_init.ExperimentType.GAUSS3D
+            if experiment == definitions.Experiments.GAUSS3D
+            else driver_init.ExperimentType.JABW
+        )
         config = icon4py_configuration.read_config(
-            experiment_type=driver_init.ExperimentType.GAUSS3D,
+            experiment_type=experiment_type,
             backend=backend,
         )
         diffusion_config = config.diffusion_config
@@ -351,7 +369,7 @@ def test_run_timeloop_single_step(
     assert test_utils.dallclose(
         prognostic_states.current.w.asnumpy(),
         w_sp.asnumpy(),
-        atol=8e-14,
+        atol=9e-14,
     )
 
     assert test_utils.dallclose(

model/standalone_driver/pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
   "icon4py-testing>=0.0.6",  # TODO(): remove this dependency when driver is fully standalone
   # external dependencies
   "typer>=0.20.0",
+  "pytest",
   "devtools>=0.12",
   "gt4py==1.1.7",
   "packaging>=20.0",

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

@@ -147,10 +147,10 @@ def create_static_field_factories(
         metadata=metrics_attributes.attrs,
         rayleigh_type=constants.RayleighType.KLEMP,
         rayleigh_coeff=0.1,
-        exner_expol=0.333,
-        vwind_offctr=0.2,
-        thslp_zdiffu=0.02,
-        thhgtd_zdiffu=125.0,
+        exner_expol=1.0 / 3.0,
+        vwind_offctr=0.15,
+        thslp_zdiffu=0.025,
+        thhgtd_zdiffu=200.0,
     )
 
     return driver_states.StaticFieldFactories(
@@ -579,7 +579,9 @@ def configure_logging(
     display_icon4py_logo_in_log_file()
 
 
-def get_backend_from_name(backend_name: str) -> model_backends.BackendLike:
+def get_backend_from_name(
+    backend_name: str,
+) -> model_backends.BackendLike:
     if backend_name not in model_backends.BACKENDS:
         raise ValueError(
             f"Invalid driver backend: {backend_name}. \n"