Combine compute_perturbed_quantities_and_interpolation further

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

@@ -22,7 +22,7 @@
 from typing import Final
 
 import gt4py.next as gtx
-from gt4py.next import astype, broadcast, maximum
+from gt4py.next import astype, maximum
 from gt4py.next.experimental import concat_where
 
 from icon4py.model.atmosphere.dycore.dycore_states import HorizontalPressureDiscretizationType
@@ -32,12 +32,6 @@
 from icon4py.model.atmosphere.dycore.stencils.extrapolate_temporally_exner_pressure import (
     _extrapolate_temporally_exner_pressure,
 )
-from icon4py.model.atmosphere.dycore.stencils.init_cell_kdim_field_with_zero_wp import (
-    _init_cell_kdim_field_with_zero_wp,
-)
-from icon4py.model.atmosphere.dycore.stencils.init_two_cell_kdim_fields_with_zero_vp import (
-    _init_two_cell_kdim_fields_with_zero_vp,
-)
 from icon4py.model.atmosphere.dycore.stencils.interpolate_to_surface import _interpolate_to_surface
 from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
 from icon4py.model.common.dimension import Koff
@@ -74,6 +68,9 @@ def _calculate_pressure_buoyancy_acceleration_at_cells_on_half_levels(
 
 @gtx.field_operator
 def _compute_perturbed_quantities_and_interpolation(
+    time_extrapolation_parameter_for_exner: fa.CellKField[ta.vpfloat],
+    current_exner: fa.CellKField[ta.wpfloat],
+    reference_exner_at_cells_on_model_levels: fa.CellKField[ta.vpfloat],
     current_rho: fa.CellKField[ta.wpfloat],
     reference_rho_at_cells_on_model_levels: fa.CellKField[ta.wpfloat],
     current_theta_v: fa.CellKField[ta.wpfloat],
@@ -86,27 +83,41 @@ def _compute_perturbed_quantities_and_interpolation(
     pressure_buoyancy_acceleration_at_cells_on_half_levels: fa.CellKField[ta.vpfloat],
     rho_at_cells_on_half_levels: fa.CellKField[ta.wpfloat],
     exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat],
-    temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat],
-    theta_v_at_cells_on_half_levels: fa.CellKField[ta.wpfloat],
+    wgtfacq_c: fa.CellKField[ta.vpfloat],
+    reference_theta_at_cells_on_half_levels: fa.CellKField[ta.vpfloat],
     igradp_method: gtx.int32,
     nflatlev: gtx.int32,
+    surface_level: gtx.int32,
 ) -> tuple[
-    fa.CellKField[ta.vpfloat],
-    fa.CellKField[ta.vpfloat],
-    fa.CellKField[ta.vpfloat],
     fa.CellKField[ta.wpfloat],
     fa.CellKField[ta.wpfloat],
     fa.CellKField[ta.vpfloat],
     fa.CellKField[ta.vpfloat],
     fa.CellKField[ta.wpfloat],
+    fa.CellKField[ta.wpfloat],
+    fa.CellKField[ta.wpfloat],
+    fa.CellKField[ta.wpfloat],
+    fa.CellKField[ta.wpfloat],
 ]:
+    (
+        temporal_extrapolation_of_perturbed_exner,
+        perturbed_exner_at_cells_on_model_levels,
+    ) = _extrapolate_temporally_exner_pressure(
+        exner_exfac=time_extrapolation_parameter_for_exner,
+        exner=current_exner,
+        exner_ref_mc=reference_exner_at_cells_on_model_levels,
+        exner_pr=perturbed_exner_at_cells_on_model_levels,
+    )
+
     exner_at_cells_on_half_levels = (
         concat_where(
-            (maximum(1, nflatlev) <= dims.KDim),
+            dims.KDim < surface_level - 1,
             _interpolate_cell_field_to_half_levels_vp(
                 wgtfac_c=wgtfac_c, interpolant=temporal_extrapolation_of_perturbed_exner
             ),
-            exner_at_cells_on_half_levels,
+            _interpolate_to_surface(
+                wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner
+            ),
         )
         if igradp_method == horzpres_discr_type.TAYLOR_HYDRO
         else exner_at_cells_on_half_levels
@@ -131,19 +142,22 @@ def _compute_perturbed_quantities_and_interpolation(
     wgtfac_c_wp = astype(wgtfac_c, wpfloat)
 
     perturbed_theta_v_at_cells_on_half_levels = concat_where(
-        dims.KDim >= 1,
+        dims.KDim < surface_level - 1,
         _interpolate_cell_field_to_half_levels_vp(
             wgtfac_c=wgtfac_c, interpolant=perturbed_theta_v_at_cells_on_model_levels
         ),
-        broadcast(0.0, (dims.CellDim, dims.KDim)),
+        _interpolate_to_surface(
+            wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels
+        ),
     )
 
     theta_v_at_cells_on_half_levels = concat_where(
-        dims.KDim >= 1,
+        dims.KDim < surface_level - 1,
         _interpolate_cell_field_to_half_levels_wp(
             wgtfac_c=wgtfac_c_wp, interpolant=current_theta_v
         ),
-        theta_v_at_cells_on_half_levels,
+        reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels,
+        # theta_v_at_cells_on_half_levels,
     )
 
     ddqz_z_half_wp = astype(ddqz_z_half, wpfloat)
@@ -170,30 +184,31 @@ def _compute_perturbed_quantities_and_interpolation(
         perturbed_theta_v_at_cells_on_half_levels,
         theta_v_at_cells_on_half_levels,
         pressure_buoyancy_acceleration_at_cells_on_half_levels,
+        temporal_extrapolation_of_perturbed_exner,
     )
 
 
 @gtx.field_operator
-def _surface_computations(
-    wgtfacq_c: fa.CellKField[ta.wpfloat],
-    exner_at_cells_on_half_levels: fa.CellKField[ta.vpfloat],
-    igradp_method: gtx.int32,
-) -> tuple[
-    fa.CellKField[ta.vpfloat],
-    fa.CellKField[ta.vpfloat],
-]:
-    temporal_extrapolation_of_perturbed_exner = _init_cell_kdim_field_with_zero_wp()
+def _set_theta_v_and_exner_on_surface_level(
+    temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat],
+    wgtfacq_c: fa.CellKField[vpfloat],
+    perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat],
+    reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat],
+) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat], fa.CellKField[vpfloat]]:
+    perturbed_theta_v_at_cells_on_half_levels = _interpolate_to_surface(
+        wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels
+    )
+    theta_v_at_cells_on_half_levels = (
+        reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels
+    )
 
-    exner_at_cells_on_half_levels = (
-        _interpolate_to_surface(
-            wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner
-        )
-        if igradp_method == horzpres_discr_type.TAYLOR_HYDRO
-        else exner_at_cells_on_half_levels
+    exner_at_cells_on_half_levels = _interpolate_to_surface(
+        wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner
     )
 
     return (
-        temporal_extrapolation_of_perturbed_exner,
+        perturbed_theta_v_at_cells_on_half_levels,
+        astype(theta_v_at_cells_on_half_levels, wpfloat),
         exner_at_cells_on_half_levels,
     )
 
@@ -217,7 +232,7 @@ def _compute_first_and_second_vertical_derivative_of_exner(
 ]:
     ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = (
         concat_where(
-            (nflatlev <= dims.KDim),
+            nflatlev <= dims.KDim,
             _compute_first_vertical_derivative_at_cells(
                 exner_at_cells_on_half_levels, inv_ddqz_z_full
             ),
@@ -229,7 +244,7 @@ def _compute_first_and_second_vertical_derivative_of_exner(
 
     d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = (
         concat_where(
-            (nflat_gradp <= dims.KDim),
+            nflat_gradp <= dims.KDim,
             -vpfloat("0.5")
             * (
                 (
@@ -251,31 +266,6 @@ def _compute_first_and_second_vertical_derivative_of_exner(
     )
 
 
-@gtx.field_operator
-def _set_theta_v_and_exner_on_surface_level(
-    temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat],
-    wgtfacq_c: fa.CellKField[vpfloat],
-    perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat],
-    reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat],
-) -> tuple[fa.CellKField[vpfloat], fa.CellKField[wpfloat], fa.CellKField[vpfloat]]:
-    perturbed_theta_v_at_cells_on_half_levels = _interpolate_to_surface(
-        wgtfacq_c=wgtfacq_c, interpolant=perturbed_theta_v_at_cells_on_model_levels
-    )
-    theta_v_at_cells_on_half_levels = (
-        reference_theta_at_cells_on_half_levels + perturbed_theta_v_at_cells_on_half_levels
-    )
-
-    exner_at_cells_on_half_levels = _interpolate_to_surface(
-        wgtfacq_c=wgtfacq_c, interpolant=temporal_extrapolation_of_perturbed_exner
-    )
-
-    return (
-        perturbed_theta_v_at_cells_on_half_levels,
-        astype(theta_v_at_cells_on_half_levels, wpfloat),
-        exner_at_cells_on_half_levels,
-    )
-
-
 @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
 def compute_perturbed_quantities_and_interpolation(
     temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat],
@@ -374,41 +364,11 @@ def compute_perturbed_quantities_and_interpolation(
         - ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_level
         - d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels
     """
-    _init_two_cell_kdim_fields_with_zero_vp(
-        out=(perturbed_rho_at_cells_on_model_levels, perturbed_theta_v_at_cells_on_model_levels),
-        domain={
-            dims.CellDim: (start_cell_lateral_boundary, start_cell_lateral_boundary_level_3),
-            dims.KDim: (model_top, surface_level - 1),
-        },
-    )
-
-    _extrapolate_temporally_exner_pressure(
-        exner_exfac=time_extrapolation_parameter_for_exner,
-        exner=current_exner,
-        exner_ref_mc=reference_exner_at_cells_on_model_levels,
-        exner_pr=perturbed_exner_at_cells_on_model_levels,
-        out=(temporal_extrapolation_of_perturbed_exner, perturbed_exner_at_cells_on_model_levels),
-        domain={
-            dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
-            dims.KDim: (model_top, surface_level - 1),
-        },
-    )
-
-    _surface_computations(
-        wgtfacq_c=wgtfacq_c,
-        exner_at_cells_on_half_levels=exner_at_cells_on_half_levels,
-        igradp_method=igradp_method,
-        out=(
-            temporal_extrapolation_of_perturbed_exner,
-            exner_at_cells_on_half_levels,
-        ),
-        domain={
-            dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
-            dims.KDim: (surface_level - 1, surface_level),
-        },
-    )
 
     _compute_perturbed_quantities_and_interpolation(
+        time_extrapolation_parameter_for_exner=time_extrapolation_parameter_for_exner,
+        current_exner=current_exner,
+        reference_exner_at_cells_on_model_levels=reference_exner_at_cells_on_model_levels,
         current_rho=current_rho,
         reference_rho_at_cells_on_model_levels=reference_rho_at_cells_on_model_levels,
         current_theta_v=current_theta_v,
@@ -421,10 +381,11 @@ def compute_perturbed_quantities_and_interpolation(
         pressure_buoyancy_acceleration_at_cells_on_half_levels=pressure_buoyancy_acceleration_at_cells_on_half_levels,
         rho_at_cells_on_half_levels=rho_at_cells_on_half_levels,
         exner_at_cells_on_half_levels=exner_at_cells_on_half_levels,
-        temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner,
-        theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels,
+        wgtfacq_c=wgtfacq_c,
+        reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels,
         igradp_method=igradp_method,
         nflatlev=nflatlev,
+        surface_level=surface_level,
         out=(
             perturbed_rho_at_cells_on_model_levels,
             perturbed_theta_v_at_cells_on_model_levels,
@@ -434,27 +395,46 @@ def compute_perturbed_quantities_and_interpolation(
             perturbed_theta_v_at_cells_on_half_levels,
             theta_v_at_cells_on_half_levels,
             pressure_buoyancy_acceleration_at_cells_on_half_levels,
+            temporal_extrapolation_of_perturbed_exner,
         ),
-        domain={
-            dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
-            dims.KDim: (model_top, surface_level - 1),
-        },
-    )
-
-    _set_theta_v_and_exner_on_surface_level(
-        temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner,
-        wgtfacq_c=wgtfacq_c,
-        perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels,
-        reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels,
-        out=(
-            perturbed_theta_v_at_cells_on_half_levels,
-            theta_v_at_cells_on_half_levels,
-            exner_at_cells_on_half_levels,
+        domain=(
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (model_top, surface_level - 1),
+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (model_top, surface_level - 1),
+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (model_top, surface_level - 1),
+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (model_top, surface_level - 1),
+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (maximum(1, nflatlev), surface_level),
+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (1, surface_level),
+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (1, surface_level),
+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (model_top, surface_level - 1),
+            },
+            {
+                dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
+                dims.KDim: (model_top, surface_level - 1),
+            },
         ),
-        domain={
-            dims.CellDim: (start_cell_lateral_boundary_level_3, end_cell_halo),
-            dims.KDim: (surface_level - 1, surface_level),
-        },
     )
 
     _compute_first_and_second_vertical_derivative_of_exner(