undo fusing of _set_theta_v_and_exner_on_surface_level

Author: edopao

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

@@ -68,9 +68,6 @@ 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],
@@ -83,6 +80,7 @@ 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,
@@ -95,18 +93,7 @@ def _compute_perturbed_quantities_and_interpolation(
     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,
@@ -177,32 +164,6 @@ 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 _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,
     )
 
 
@@ -216,43 +177,16 @@ 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) & (dims.KDim < surface_level - 1),
+            nflatlev <= dims.KDim,
             _compute_first_vertical_derivative_at_cells(
                 exner_at_cells_on_half_levels, inv_ddqz_z_full
             ),
@@ -264,7 +198,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) & (dims.KDim < surface_level - 1),
+            nflat_gradp <= dims.KDim,
             -vpfloat("0.5")
             * (
                 (
@@ -283,8 +217,30 @@ 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,
-        theta_v_at_cells_on_half_levels,
+        astype(theta_v_at_cells_on_half_levels, wpfloat),
         exner_at_cells_on_half_levels,
     )
 
@@ -388,10 +344,22 @@ def compute_perturbed_quantities_and_interpolation(
         - d2dz2_of_temporal_extrapolation_of_perturbed_exner_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,
+        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),
+        },
+    )
+
     _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,
@@ -404,6 +372,7 @@ 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,
@@ -416,14 +385,29 @@ 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,
@@ -433,24 +417,16 @@ 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),
+            dims.KDim: (model_top, surface_level - 1),
         },
     )