add init_exner_pr stencil (#461)

* add init_exner_pr stencil and docstring to initialization functions, and some clean up.

* remove a bug in solve_nonhydro corrector stencil 41

* add options for grid root and grid level

* add C2E2C2EDim size in IconSerialDataProvider

* testing

* Correct docstring of jabw wind computation and remove c2e2c2e grid size

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/init_exner_pr.py

@@ -0,0 +1,50 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+from gt4py.next.common import GridType
+from gt4py.next.ffront.decorator import field_operator, program
+from gt4py.next.ffront.fbuiltins import Field, int32
+
+from icon4py.model.common.dimension import CellDim, KDim
+from icon4py.model.common.settings import backend
+from icon4py.model.common.type_alias import vpfloat
+
+
+@field_operator
+def _init_exner_pr(
+    exner: Field[[CellDim, KDim], vpfloat],
+    exner_ref: Field[[CellDim, KDim], vpfloat],
+) -> Field[[CellDim, KDim], vpfloat]:
+    exner_pr = exner - exner_ref
+    return exner_pr
+
+
+@program(grid_type=GridType.UNSTRUCTURED, backend=backend)
+def init_exner_pr(
+    exner: Field[[CellDim, KDim], vpfloat],
+    exner_ref: Field[[CellDim, KDim], vpfloat],
+    exner_pr: Field[[CellDim, KDim], vpfloat],
+    horizontal_start: int32,
+    horizontal_end: int32,
+    vertical_start: int32,
+    vertical_end: int32,
+):
+    _init_exner_pr(
+        exner,
+        exner_ref,
+        out=exner_pr,
+        domain={
+            CellDim: (horizontal_start, horizontal_end),
+            KDim: (vertical_start, vertical_end),
+        },
+    )

model/atmosphere/dycore/src/icon4py/model/atmosphere/dycore/nh_solve/solve_nonhydro.py

@@ -1647,20 +1647,20 @@ def run_corrector_step(
                 offset_provider={},
             )
 
-            # verified for e-9
-            log.debug(f"corrector: start stencile 41")
-            compute_divergence_of_fluxes_of_rho_and_theta(
-                geofac_div=self.interpolation_state.geofac_div,
-                mass_fl_e=diagnostic_state_nh.mass_fl_e,
-                z_theta_v_fl_e=self.z_theta_v_fl_e,
-                z_flxdiv_mass=self.z_flxdiv_mass,
-                z_flxdiv_theta=self.z_flxdiv_theta,
-                horizontal_start=start_cell_nudging,
-                horizontal_end=end_cell_local,
-                vertical_start=0,
-                vertical_end=self.grid.num_levels,
-                offset_provider=self.grid.offset_providers,
-            )
+        # verified for e-9
+        log.debug(f"corrector: start stencil 41")
+        compute_divergence_of_fluxes_of_rho_and_theta(
+            geofac_div=self.interpolation_state.geofac_div,
+            mass_fl_e=diagnostic_state_nh.mass_fl_e,
+            z_theta_v_fl_e=self.z_theta_v_fl_e,
+            z_flxdiv_mass=self.z_flxdiv_mass,
+            z_flxdiv_theta=self.z_flxdiv_theta,
+            horizontal_start=start_cell_nudging,
+            horizontal_end=end_cell_local,
+            vertical_start=0,
+            vertical_end=self.grid.num_levels,
+            offset_provider=self.grid.offset_providers,
+        )
 
         if self.config.itime_scheme == 4:
             log.debug(f"corrector start stencil 42 44 45 45b")

model/atmosphere/dycore/tests/dycore_stencil_tests/test_init_exner_pr.py

@@ -0,0 +1,51 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+import pytest
+from gt4py.next.ffront.fbuiltins import int32
+
+from icon4py.model.atmosphere.dycore.init_exner_pr import (
+    init_exner_pr,
+)
+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.type_alias import vpfloat
+
+
+class TestInitExnerPr(StencilTest):
+    PROGRAM = init_exner_pr
+    OUTPUTS = ("exner_pr",)
+
+    @staticmethod
+    def reference(grid, exner: np.array, exner_ref: np.array, **kwargs) -> dict:
+        exner_pr = exner - exner_ref
+        return dict(
+            exner_pr=exner_pr,
+        )
+
+    @pytest.fixture
+    def input_data(self, grid):
+        exner = random_field(grid, CellDim, KDim, dtype=vpfloat)
+        exner_ref = random_field(grid, CellDim, KDim, dtype=vpfloat)
+        exner_pr = zero_field(grid, CellDim, KDim, dtype=vpfloat)
+
+        return dict(
+            exner=exner,
+            exner_ref=exner_ref,
+            exner_pr=exner_pr,
+            horizontal_start=int32(0),
+            horizontal_end=int32(grid.num_cells),
+            vertical_start=int32(0),
+            vertical_end=int32(grid.num_levels),
+        )

model/common/src/icon4py/model/common/test_utils/serialbox_utils.py

@@ -536,12 +536,14 @@ def rbf_vec_coeff_e(self):
         ).transpose()
         return as_field((EdgeDim, E2C2EDim), buffer)
 
+    @IconSavepoint.optionally_registered()
     def rbf_vec_coeff_c1(self):
         buffer = np.squeeze(
             self.serializer.read("rbf_vec_coeff_c1", self.savepoint).astype(float)
         ).transpose()
         return as_field((CellDim, C2E2C2EDim), buffer)
 
+    @IconSavepoint.optionally_registered()
     def rbf_vec_coeff_c2(self):
         buffer = np.squeeze(
             self.serializer.read("rbf_vec_coeff_c2", self.savepoint).astype(float)
@@ -580,6 +582,7 @@ def hmask_dd3d(self):
     def inv_ddqz_z_full(self):
         return self._get_field("inv_ddqz_z_full", CellDim, KDim)
 
+    @IconSavepoint.optionally_registered(CellDim, KDim)
     def ddqz_z_full(self):
         return self._get_field("ddqz_z_full", CellDim, KDim)
 

model/driver/src/icon4py/model/driver/dycore_driver.py

@@ -186,10 +186,10 @@ def time_integration(
         for time_step in range(self._n_time_steps):
             log.info(f"simulation date : {self._simulation_date} run timestep : {time_step}")
             log.info(
-                f" MAX VN: {prognostic_state_list[self._now].vn.asnumpy().max():.5e} , MAX W: {prognostic_state_list[self._now].w.asnumpy().max():.5e}"
+                f" MAX VN: {prognostic_state_list[self._now].vn.asnumpy().max():.15e} , MAX W: {prognostic_state_list[self._now].w.asnumpy().max():.15e}"
             )
             log.info(
-                f" MAX RHO: {prognostic_state_list[self._now].rho.asnumpy().max():.5e} , MAX THETA_V: {prognostic_state_list[self._now].theta_v.asnumpy().max():.5e}"
+                f" MAX RHO: {prognostic_state_list[self._now].rho.asnumpy().max():.15e} , MAX THETA_V: {prognostic_state_list[self._now].theta_v.asnumpy().max():.15e}"
             )
             # TODO (Chia Rui): check with Anurag about printing of max and min of variables.
 
@@ -294,6 +294,8 @@ def initialize(
     props: ProcessProperties,
     serialization_type: SerializationType,
     experiment_type: ExperimentType,
+    grid_root,
+    grid_level,
 ):
     """
     Inititalize the driver run.
@@ -323,16 +325,24 @@ def initialize(
     log.info(f"reading configuration: experiment {experiment_type}")
     config = read_config(experiment_type)
 
-    decomp_info = read_decomp_info(file_path, props, serialization_type)
+    decomp_info = read_decomp_info(file_path, props, serialization_type, grid_root, grid_level)
 
     log.info(f"initializing the grid from '{file_path}'")
-    icon_grid = read_icon_grid(file_path, rank=props.rank, ser_type=serialization_type)
+    icon_grid = read_icon_grid(
+        file_path,
+        rank=props.rank,
+        ser_type=serialization_type,
+        grid_root=grid_root,
+        grid_level=grid_level,
+    )
     log.info(f"reading input fields from '{file_path}'")
     (edge_geometry, cell_geometry, vertical_geometry, c_owner_mask) = read_geometry_fields(
         file_path,
         damping_height=config.run_config.damping_height,
         rank=props.rank,
         ser_type=serialization_type,
+        grid_root=grid_root,
+        grid_level=grid_level,
     )
     (
         diffusion_metric_state,
@@ -341,7 +351,11 @@ def initialize(
         solve_nonhydro_interpolation_state,
         diagnostic_metric_state,
     ) = read_static_fields(
-        file_path, rank=props.rank, ser_type=serialization_type, experiment_type=experiment_type
+        file_path,
+        rank=props.rank,
+        ser_type=serialization_type,
+        grid_root=grid_root,
+        grid_level=grid_level,
     )
 
     log.info("initializing diffusion")
@@ -418,7 +432,9 @@ def initialize(
     help="serialization type for grid info and static fields",
 )
 @click.option("--experiment_type", default="any", help="experiment selection")
-def main(input_path, run_path, mpi, serialization_type, experiment_type):
+@click.option("--grid_root", default=2, help="experiment selection")
+@click.option("--grid_level", default=4, help="experiment selection")
+def main(input_path, run_path, mpi, serialization_type, experiment_type, grid_root, grid_level):
     """
     Run the driver.
 
@@ -449,7 +465,9 @@ def main(input_path, run_path, mpi, serialization_type, experiment_type):
         diagnostic_state,
         prep_adv,
         inital_divdamp_fac_o2,
-    ) = initialize(Path(input_path), parallel_props, serialization_type, experiment_type)
+    ) = initialize(
+        Path(input_path), parallel_props, serialization_type, experiment_type, grid_root, grid_level
+    )
     log.info(f"Starting ICON dycore run: {timeloop.simulation_date.isoformat()}")
     log.info(
         f"input args: input_path={input_path}, n_time_steps={timeloop.n_time_steps}, ending date={timeloop.run_config.end_date}"