Combine compute_perturbed_quantities_and_interpolation further

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

@@ -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,23 +83,33 @@ 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],
     igradp_method: gtx.int32,
     nflatlev: 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),
+            maximum(1, nflatlev) <= dims.KDim,
             _interpolate_cell_field_to_half_levels_vp(
                 wgtfac_c=wgtfac_c, interpolant=temporal_extrapolation_of_perturbed_exner
             ),
@@ -170,30 +177,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,
     )
 
@@ -208,16 +216,43 @@ def _compute_first_and_second_vertical_derivative_of_exner(
     d2dexdz2_fac1_mc: fa.CellKField[vpfloat],
     d2dexdz2_fac2_mc: fa.CellKField[vpfloat],
     perturbed_theta_v_at_cells_on_model_levels: fa.CellKField[vpfloat],
+    temporal_extrapolation_of_perturbed_exner: fa.CellKField[vpfloat],
+    wgtfacq_c: fa.CellKField[vpfloat],
+    reference_theta_at_cells_on_half_levels: fa.CellKField[vpfloat],
+    theta_v_at_cells_on_half_levels: fa.CellKField[vpfloat],
     igradp_method: gtx.int32,
     nflatlev: gtx.int32,
     nflat_gradp: gtx.int32,
+    surface_level: gtx.int32,
 ) -> tuple[
     fa.CellKField[vpfloat],
     fa.CellKField[vpfloat],
+    fa.CellKField[vpfloat],
+    fa.CellKField[vpfloat],
+    fa.CellKField[vpfloat],
 ]:
+    (
+        perturbed_theta_v_at_cells_on_half_levels,
+        theta_v_at_cells_on_half_levels,
+        exner_at_cells_on_half_levels,
+    ) = concat_where(
+        dims.KDim >= 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,
+        ),
+        (
+            perturbed_theta_v_at_cells_on_half_levels,
+            theta_v_at_cells_on_half_levels,
+            exner_at_cells_on_half_levels,
+        ),
+    )
+
     ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels = (
         concat_where(
-            (nflatlev <= dims.KDim),
+            (nflatlev <= dims.KDim) & (dims.KDim < surface_level - 1),
             _compute_first_vertical_derivative_at_cells(
                 exner_at_cells_on_half_levels, inv_ddqz_z_full
             ),
@@ -229,7 +264,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) & (dims.KDim < surface_level - 1),
             -vpfloat("0.5")
             * (
                 (
@@ -248,30 +283,8 @@ def _compute_first_and_second_vertical_derivative_of_exner(
     return (
         ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels,
         d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels,
-    )
-
-
-@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),
+        theta_v_at_cells_on_half_levels,
         exner_at_cells_on_half_levels,
     )
 
@@ -374,41 +387,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,7 +404,6 @@ 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,
         igradp_method=igradp_method,
         nflatlev=nflatlev,
@@ -434,29 +416,14 @@ 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: (surface_level - 1, surface_level),
-        },
-    )
-
     _compute_first_and_second_vertical_derivative_of_exner(
         exner_at_cells_on_half_levels=exner_at_cells_on_half_levels,
         inv_ddqz_z_full=inv_ddqz_z_full,
@@ -466,16 +433,24 @@ def compute_perturbed_quantities_and_interpolation(
         d2dexdz2_fac1_mc=d2dexdz2_fac1_mc,
         d2dexdz2_fac2_mc=d2dexdz2_fac2_mc,
         perturbed_theta_v_at_cells_on_model_levels=perturbed_theta_v_at_cells_on_model_levels,
+        temporal_extrapolation_of_perturbed_exner=temporal_extrapolation_of_perturbed_exner,
+        wgtfacq_c=wgtfacq_c,
+        reference_theta_at_cells_on_half_levels=reference_theta_at_cells_on_half_levels,
+        theta_v_at_cells_on_half_levels=theta_v_at_cells_on_half_levels,
         igradp_method=igradp_method,
         nflatlev=nflatlev,
         nflat_gradp=nflat_gradp,
+        surface_level=surface_level,
         out=(
             ddz_of_temporal_extrapolation_of_perturbed_exner_on_model_levels,
             d2dz2_of_temporal_extrapolation_of_perturbed_exner_on_model_levels,
+            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.KDim: (model_top, surface_level),
         },
     )