Replace pg_edgeidx_dsl mask by multiplying with 0.0

model/atmosphere/dycore/docs/scidoc/icon4py.model.atmosphere.dycore.solve_nonhydro.SolveNonhydro.run_predictor_step.rst

@@ -176,7 +176,6 @@ Inputs:
  - $\exnerprimegradh{\ntilde}{\e}{\k}$ : z_gradh_exner
  - $\exnhydrocorr{\e}$ : hydro_corr_horizontal
  - $(h_k - h_{k^*})$ : pg_exdist
- - $\IDXpg$ : ipeidx_dsl
 
 scidoc:
 Outputs:

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

@@ -275,11 +275,15 @@ class MetricStateNonHydro:
     vertoffset_gradp: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim, dims.KDim], gtx.int32]
     zdiff_gradp: gtx.Field[gtx.Dims[dims.EdgeDim, dims.E2CDim, dims.KDim], ta.vpfloat]
     nflat_gradp: gtx.int32
-    """The minimum height index at which the height of the center of an edge lies within two neighboring cells so that
+    """
+    The minimum height index at which the height of the center of an edge lies within two neighboring cells so that
     horizontal pressure gradient can be computed by first order discretization scheme.
     """
-    pg_edgeidx_dsl: fa.EdgeKField[bool]
+
     pg_exdist: fa.EdgeKField[ta.vpfloat]
+    """
+    Extrapolation distance needed for HorizontalPressureDiscretizationType.TAYLOR_HYDRO.
+    """
 
     exner_w_explicit_weight_parameter: fa.CellField[ta.wpfloat]
     """

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

@@ -456,7 +456,6 @@ def __init__(
                 "c_lin_e": self._interpolation_state.c_lin_e,
                 "ikoffset": self._metric_state_nonhydro.vertoffset_gradp,
                 "zdiff_gradp": self._metric_state_nonhydro.zdiff_gradp,
-                "ipeidx_dsl": self._metric_state_nonhydro.pg_edgeidx_dsl,
                 "pg_exdist": self._metric_state_nonhydro.pg_exdist,
                 "inv_dual_edge_length": self._edge_geometry.inverse_dual_edge_lengths,
                 "iau_wgt_dyn": self._config.iau_wgt_dyn,

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

@@ -9,7 +9,7 @@
 
 import gt4py.next as gtx
 from gt4py.eve import utils as eve_utils
-from gt4py.next import broadcast
+from gt4py.next import broadcast, where
 from gt4py.next.experimental import concat_where
 
 from icon4py.model.atmosphere.dycore.stencils.add_analysis_increments_to_vn import (
@@ -30,9 +30,6 @@
 from icon4py.model.atmosphere.dycore.stencils.apply_4th_order_divergence_damping import (
     _apply_4th_order_divergence_damping,
 )
-from icon4py.model.atmosphere.dycore.stencils.apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure import (
-    _apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure,
-)
 from icon4py.model.atmosphere.dycore.stencils.apply_weighted_2nd_and_4th_order_divergence_damping import (
     _apply_weighted_2nd_and_4th_order_divergence_damping,
 )
@@ -79,6 +76,20 @@ def apply_on_vertical_level(
     )
 
 
+@gtx.field_operator
+def apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure(
+    pg_exdist: fa.EdgeKField[ta.vpfloat],
+    z_hydro_corr: fa.EdgeField[ta.wpfloat],
+    z_gradh_exner: fa.EdgeKField[ta.vpfloat],
+) -> fa.EdgeKField[ta.vpfloat]:
+    # Note: In the original Fortran code `pg_exdist` is implemented as a list,
+    # in ICON4Py it's a full field intialized with zeros for points that are not in the list.
+    z_gradh_exner_vp = where(
+        pg_exdist != 0.0, z_gradh_exner + z_hydro_corr * pg_exdist, z_gradh_exner
+    )
+    return z_gradh_exner_vp
+
+
 @gtx.field_operator
 def _compute_horizontal_pressure_gradient(
     temporal_extrapolation_of_perturbed_exner: fa.CellKField[ta.vpfloat],
@@ -89,7 +100,6 @@ def _compute_horizontal_pressure_gradient(
     c_lin_e: gtx.Field[[dims.EdgeDim, dims.E2CDim], ta.wpfloat],
     ikoffset: gtx.Field[[dims.EdgeDim, dims.E2CDim, dims.KDim], gtx.int32],
     zdiff_gradp: gtx.Field[[dims.EdgeDim, dims.E2CDim, dims.KDim], ta.vpfloat],
-    ipeidx_dsl: fa.EdgeKField[bool],
     pg_exdist: fa.EdgeKField[ta.vpfloat],
     inv_dual_edge_length: fa.EdgeField[ta.wpfloat],
     nflatlev: gtx.int32,
@@ -120,8 +130,7 @@ def _compute_horizontal_pressure_gradient(
         ),
     )
 
-    return _apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure(
-        ipeidx_dsl=ipeidx_dsl,
+    return apply_hydrostatic_correction_to_horizontal_gradient_of_exner_pressure(
         pg_exdist=pg_exdist,
         z_hydro_corr=hydrostatic_correction_on_lowest_level,
         z_gradh_exner=horizontal_pressure_gradient,
@@ -157,7 +166,6 @@ def _compute_rho_theta_pgrad_and_update_vn(
     c_lin_e: gtx.Field[[dims.EdgeDim, dims.E2CDim], ta.wpfloat],
     ikoffset: gtx.Field[[dims.EdgeDim, dims.E2CDim, dims.KDim], gtx.int32],
     zdiff_gradp: gtx.Field[[dims.EdgeDim, dims.E2CDim, dims.KDim], ta.vpfloat],
-    ipeidx_dsl: fa.EdgeKField[bool],
     pg_exdist: fa.EdgeKField[ta.vpfloat],
     inv_dual_edge_length: fa.EdgeField[ta.wpfloat],
     dtime: ta.wpfloat,
@@ -218,7 +226,6 @@ def _compute_rho_theta_pgrad_and_update_vn(
         c_lin_e=c_lin_e,
         ikoffset=ikoffset,
         zdiff_gradp=zdiff_gradp,
-        ipeidx_dsl=ipeidx_dsl,
         pg_exdist=pg_exdist,
         inv_dual_edge_length=inv_dual_edge_length,
         nflatlev=nflatlev,
@@ -401,7 +408,6 @@ def compute_rho_theta_pgrad_and_update_vn(
     c_lin_e: gtx.Field[[dims.EdgeDim, dims.E2CDim], ta.wpfloat],
     ikoffset: gtx.Field[[dims.EdgeDim, dims.E2CDim, dims.KDim], gtx.int32],
     zdiff_gradp: gtx.Field[[dims.EdgeDim, dims.E2CDim, dims.KDim], ta.vpfloat],
-    ipeidx_dsl: fa.EdgeKField[bool],
     pg_exdist: fa.EdgeKField[ta.vpfloat],
     inv_dual_edge_length: fa.EdgeField[ta.wpfloat],
     dtime: ta.wpfloat,
@@ -458,7 +464,6 @@ def compute_rho_theta_pgrad_and_update_vn(
         - c_lin_e: interpolation coefficient for computation of interpolating a cell-based variables to an edge-based variable
         - ikoffset: k offset index (offset from the lowest k index where the neighboring cell centers lie within the thickness of the layer) for hyrostatic correction
         - zdiff_gradp: vertical distance between current cell height and neighboring cell height for pressure gradient over multiple levels [m]
-        - ipeidx_dsl: A mask for hydrostatic correction
         - pg_exdist: vertical distance between current cell height and neighboring cell height for hydrostatic correction [m]
         - inv_dual_edge_length: inverse dual edge length [m]
         - dtime: time step [s]
@@ -508,7 +513,6 @@ def compute_rho_theta_pgrad_and_update_vn(
         c_lin_e=c_lin_e,
         ikoffset=ikoffset,
         zdiff_gradp=zdiff_gradp,
-        ipeidx_dsl=ipeidx_dsl,
         pg_exdist=pg_exdist,
         inv_dual_edge_length=inv_dual_edge_length,
         dtime=dtime,

model/atmosphere/dycore/tests/dycore/integration_tests/test_benchmark_solve_nonhydro.py

@@ -175,7 +175,6 @@ def solve_nonhydro(
         zdiff_gradp=metrics_field_source.get(metrics_attributes.ZDIFF_GRADP),
         vertoffset_gradp=metrics_field_source.get(metrics_attributes.VERTOFFSET_GRADP),
         nflat_gradp=metrics_field_source.get(metrics_attributes.NFLAT_GRADP),
-        pg_edgeidx_dsl=metrics_field_source.get(metrics_attributes.PG_EDGEIDX_DSL),
         pg_exdist=metrics_field_source.get(metrics_attributes.PG_EDGEDIST_DSL),
         ddqz_z_full_e=metrics_field_source.get(metrics_attributes.DDQZ_Z_FULL_E),
         ddxt_z_full=metrics_field_source.get(metrics_attributes.DDXT_Z_FULL),

model/atmosphere/dycore/tests/dycore/integration_tests/test_solve_nonhydro.py

@@ -1513,7 +1513,6 @@ def test_compute_rho_theta_pgrad_and_update_vn(
         c_lin_e=interpolation_savepoint.c_lin_e(),
         ikoffset=metrics_savepoint.vertoffset_gradp(),
         zdiff_gradp=metrics_savepoint.zdiff_gradp(),
-        ipeidx_dsl=metrics_savepoint.pg_edgeidx_dsl(),
         pg_exdist=metrics_savepoint.pg_exdist(),
         inv_dual_edge_length=grid_savepoint.inv_dual_edge_length(),
         dtime=savepoint_nonhydro_init.get_metadata("dtime").get("dtime"),

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

@@ -193,7 +193,6 @@ def reference(
         c_lin_e: np.ndarray,
         ikoffset: np.ndarray,
         zdiff_gradp: np.ndarray,
-        ipeidx_dsl: np.ndarray,
         pg_exdist: np.ndarray,
         inv_dual_edge_length: np.ndarray,
         dtime: ta.wpfloat,
@@ -395,10 +394,8 @@ def at_neighbor(i: int) -> np.ndarray:
             horizontal_pressure_gradient.shape[1],
             axis=1,
         )
-        horizontal_pressure_gradient = np.where(
-            ipeidx_dsl,
-            horizontal_pressure_gradient + hydrostatic_correction * pg_exdist,
-            horizontal_pressure_gradient,
+        horizontal_pressure_gradient = (
+            horizontal_pressure_gradient + hydrostatic_correction * pg_exdist
         )
 
         next_vn = np.where(
@@ -477,8 +474,9 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict:
         )
         hydrostatic_correction_on_lowest_level = data_alloc.random_field(grid, dims.EdgeDim)
         zdiff_gradp = data_alloc.random_field(grid, dims.EdgeDim, dims.E2CDim, dims.KDim)
-        ipeidx_dsl = data_alloc.random_mask(grid, dims.EdgeDim, dims.KDim)
-        pg_exdist = data_alloc.random_field(grid, dims.EdgeDim, dims.KDim)
+        pg_exdist = data_alloc.random_field(
+            grid, dims.EdgeDim, dims.KDim
+        )  # TODO(havogt): should be allocated with a sparse pattern
         inv_dual_edge_length = data_alloc.random_field(grid, dims.EdgeDim)
         predictor_normal_wind_advective_tendency = data_alloc.random_field(
             grid, dims.EdgeDim, dims.KDim
@@ -554,7 +552,6 @@ def input_data(self, request: pytest.FixtureRequest, grid: base.Grid) -> dict:
             c_lin_e=c_lin_e,
             ikoffset=ikoffset,
             zdiff_gradp=zdiff_gradp,
-            ipeidx_dsl=ipeidx_dsl,
             pg_exdist=pg_exdist,
             inv_dual_edge_length=inv_dual_edge_length,
             dtime=dtime,

model/atmosphere/dycore/tests/dycore/utils.py

@@ -67,7 +67,6 @@ def construct_metric_state(
         zdiff_gradp=metrics_savepoint.zdiff_gradp(),
         vertoffset_gradp=metrics_savepoint.vertoffset_gradp(),
         nflat_gradp=grid_savepoint.nflat_gradp(),
-        pg_edgeidx_dsl=metrics_savepoint.pg_edgeidx_dsl(),
         pg_exdist=metrics_savepoint.pg_exdist(),
         ddqz_z_full_e=metrics_savepoint.ddqz_z_full_e(),
         ddxt_z_full=metrics_savepoint.ddxt_z_full(),

model/common/src/icon4py/model/common/metrics/metric_fields.py

@@ -635,7 +635,7 @@ def _compute_pressure_gradient_downward_extrapolation_mask_distance(
     k_lev: fa.KField[gtx.int32],
     horizontal_start_distance: int32,
     horizontal_end_distance: int32,
-) -> tuple[fa.EdgeKField[bool], fa.EdgeKField[wpfloat]]:
+) -> fa.EdgeKField[wpfloat]:
     """
     Compute an edge mask and extrapolation distance for grid points requiring downward extrapolation of the pressure gradient.
 
@@ -653,7 +653,6 @@ def _compute_pressure_gradient_downward_extrapolation_mask_distance(
         horizontal_end_distance: end index in edge fields until where extrapolation distance is computed
 
     Returns:
-        pg_edge_mask: edge index mask for points requiring downward extrapolation
         pg_exdist_dsl: extrapolation distance
 
     """
@@ -667,17 +666,14 @@ def _compute_pressure_gradient_downward_extrapolation_mask_distance(
         downward_distance,
         0.0,
     )
-    flatness_condition = (k_lev >= (flat_idx_max + 1)) & (z_me < downward_distance) & e_owner_mask
-    pg_edgeidx, pg_vertidx = where(flatness_condition, (e_lev, k_lev), (0, 0))
-    pg_edge_mask = (pg_edgeidx > 0) & (pg_vertidx > 0)
 
     pg_exdist_dsl = where(
         (k_lev >= (flat_idx_max + 1)) & (z_me < extrapolation_distance) & e_owner_mask,
         z_me - extrapolation_distance,
         0.0,
     )
 
-    return pg_edge_mask, pg_exdist_dsl
+    return pg_exdist_dsl
 
 
 @gtx.program(grid_type=gtx.GridType.UNSTRUCTURED)
@@ -689,7 +685,6 @@ def compute_pressure_gradient_downward_extrapolation_mask_distance(
     flat_idx_max: fa.EdgeField[gtx.int32],
     e_lev: fa.EdgeField[gtx.int32],
     k_lev: fa.KField[gtx.int32],
-    pg_edgeidx_dsl: fa.EdgeKField[bool],
     pg_exdist_dsl: fa.EdgeKField[wpfloat],
     horizontal_start_distance: int32,
     horizontal_end_distance: int32,
@@ -708,7 +703,7 @@ def compute_pressure_gradient_downward_extrapolation_mask_distance(
         k_lev=k_lev,
         horizontal_start_distance=horizontal_start_distance,
         horizontal_end_distance=horizontal_end_distance,
-        out=(pg_edgeidx_dsl, pg_exdist_dsl),
+        out=pg_exdist_dsl,
         domain={
             dims.EdgeDim: (horizontal_start, horizontal_end),
             dims.KDim: (vertical_start, vertical_end),

model/common/src/icon4py/model/common/metrics/metrics_attributes.py

@@ -49,7 +49,6 @@
 WGTFAC_E: Final[str] = "wgtfac_e"
 FLAT_IDX_MAX: Final[str] = "flat_idx_max"
 NFLAT_GRADP: Final[str] = "nflat_gradp"
-PG_EDGEIDX_DSL: Final[str] = "edge_mask_for_pressure_gradient_extrapolation"
 PG_EDGEDIST_DSL: Final[str] = "distance_for_pressure_gradient_extrapolation"
 MASK_PROG_HALO_C: Final[str] = "mask_prog_halo_c"
 HORIZONTAL_MASK_FOR_3D_DIVDAMP: Final[str] = "horizontal_mask_for_3d_divdamp"
@@ -302,14 +301,6 @@
         icon_var_name="flat_idx_max",
         dtype=ta.wpfloat,
     ),
-    PG_EDGEIDX_DSL: dict(
-        standard_name=PG_EDGEIDX_DSL,
-        long_name="edge mask for pressure gradient downward extrapolation",
-        units="",
-        dims=(dims.EdgeDim, dims.KDim),
-        icon_var_name="pg_edgeidx_dsl",
-        dtype=bool,
-    ),
     PG_EDGEDIST_DSL: dict(
         standard_name=PG_EDGEDIST_DSL,
         long_name="extrapolation distance for pressure gradient downward extrapolation",