Fix interpolation field geofac grg

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

@@ -32,8 +32,8 @@ def reference(grid, z_raylfac: np.array, w_1: np.array, w: np.array, **kwargs) -
     @pytest.fixture
     def input_data(self, grid):
         z_raylfac = random_field(grid, dims.KDim, dtype=wpfloat)
-        w_1 = random_field(grid, dims.CellDim, dtype=wpfloat)
         w = random_field(grid, dims.CellDim, dims.KDim, dtype=wpfloat)
+        w_1 = w[dims.KDim(0)]
 
         return dict(
             z_raylfac=z_raylfac,

model/common/src/icon4py/model/common/interpolation/interpolation_fields.py

@@ -152,9 +152,8 @@ def _compute_primal_normal_ec(
     owner_mask: data_alloc.NDArray,
     c2e: data_alloc.NDArray,
     e2c: data_alloc.NDArray,
-    horizontal_start: np.int32,
     array_ns: ModuleType = np,
-) -> data_alloc.NDArray:
+) -> tuple[data_alloc.NDArray, data_alloc.NDArray]:
     """
     Compute primal_normal_ec.
 
@@ -164,57 +163,36 @@ def _compute_primal_normal_ec(
         owner_mask: ndarray, representing a gtx.Field[gtx.Dims[CellDim], bool]
         c2e: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32]
         e2c: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], gtx.int32]
-        horizontal_start: start index to compute from
         array_ns: module - the array interface implementation to compute on, defaults to numpy
     Returns:
         primal_normal_ec: numpy array, representing a gtx.Field[gtx.Dims[CellDim, C2EDim, 2], ta.wpfloat]
     """
-    num_cells = c2e.shape[0]
-    primal_normal_ec = np.zeros([c2e.shape[0], c2e.shape[1], 2])
-    index = np.transpose(
-        np.vstack(
-            (
-                array_ns.arange(num_cells),
-                array_ns.arange(num_cells),
-                array_ns.arange(num_cells),
-            )
-        )
-    )
-    owned = np.vstack((owner_mask, owner_mask, owner_mask)).T
-    for i in range(2):
-        mask = e2c[c2e, i] == index
-        primal_normal_ec[horizontal_start:, :, 0] = primal_normal_ec[
-            horizontal_start:, :, 0
-        ] + array_ns.where(
-            owned[horizontal_start:, :],
-            mask[horizontal_start:, :] * primal_normal_cell_x[c2e[horizontal_start:], i],
-            0.0,
-        )
-        primal_normal_ec[horizontal_start:, :, 1] = primal_normal_ec[
-            horizontal_start:, :, 1
-        ] + array_ns.where(
-            owned[horizontal_start:, :],
-            mask[horizontal_start:, :] * primal_normal_cell_y[c2e[horizontal_start:], i],
-            0.0,
-        )
-    return primal_normal_ec
+
+    owned = np.stack((owner_mask, owner_mask, owner_mask)).T
+
+    inv_neighbor_index = create_inverse_neighbor_index(e2c, c2e, array_ns)
+    u_component = primal_normal_cell_x[c2e, inv_neighbor_index]
+    v_component = primal_normal_cell_y[c2e, inv_neighbor_index]
+    return (np.where(owned, u_component, 0.0), np.where(owned, v_component, 0.0))
 
 
 def _compute_geofac_grg(
-    primal_normal_ec: data_alloc.NDArray,
+    primal_normal_ec_u: data_alloc.NDArray,
+    primal_normal_ec_v: data_alloc.NDArray,
     geofac_div: data_alloc.NDArray,
     c_lin_e: data_alloc.NDArray,
     c2e: data_alloc.NDArray,
     e2c: data_alloc.NDArray,
     c2e2c: data_alloc.NDArray,
-    horizontal_start: np.int32,
+    horizontal_start: gtx.int32,
     array_ns: ModuleType = np,
 ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]:
     """
     Compute geometrical factor for Green-Gauss gradient.
 
     Args:
-        primal_normal_ec: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim, 2], ta.wpfloat]
+        primal_normal_ec_u: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim, 2], ta.wpfloat]
+        primal_normal_ec_v: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim, 2], ta.wpfloat]
         geofac_div: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], ta.wpfloat]
         c_lin_e: ndarray, representing a gtx.Field[gtx.Dims[EdgeDim, E2CDim], ta.wpfloat]
         c2e: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim], gtx.int32]
@@ -226,39 +204,29 @@ def _compute_geofac_grg(
         geofac_grg: ndarray, representing a gtx.Field[gtx.Dims[CellDim, C2EDim + 1, 2], ta.wpfloat]
     """
     num_cells = c2e.shape[0]
-    geofac_grg = array_ns.zeros([num_cells, c2e.shape[1] + 1, primal_normal_ec.shape[2]])
-    index = array_ns.transpose(
-        array_ns.vstack(
-            (
-                array_ns.arange(num_cells),
-                array_ns.arange(num_cells),
-                array_ns.arange(num_cells),
-            )
-        )
-    )
-    for k in range(e2c.shape[1]):
-        mask = e2c[c2e, k] == index
-        for i in range(primal_normal_ec.shape[2]):
-            for j in range(c2e.shape[1]):
-                geofac_grg[horizontal_start:, 0, i] = (
-                    geofac_grg[horizontal_start:, 0, i]
-                    + mask[horizontal_start:, j]
-                    * (primal_normal_ec[:, :, i] * geofac_div * c_lin_e[c2e, k])[
-                        horizontal_start:, j
-                    ]
-                )
+    targ_local_size = c2e.shape[1] + 1
+    target_shape = (num_cells, targ_local_size)
+    geofac_grg_x = array_ns.zeros(target_shape)
+    geofac_grg_y = array_ns.zeros(target_shape)
+
+    inverse_neighbor = create_inverse_neighbor_index(e2c, c2e, array_ns)
+
+    tmp = geofac_div * c_lin_e[c2e, inverse_neighbor]
+    geofac_grg_x[horizontal_start:, 0] = np.sum(primal_normal_ec_u * tmp, axis=1)[horizontal_start:]
+    geofac_grg_y[horizontal_start:, 0] = np.sum(primal_normal_ec_v * tmp, axis=1)[horizontal_start:]
+
     for k in range(e2c.shape[1]):
-        mask = e2c[c2e, k] == c2e2c
-        for i in range(primal_normal_ec.shape[2]):
-            for j in range(c2e.shape[1]):
-                geofac_grg[horizontal_start:, 1 + j, i] = (
-                    geofac_grg[horizontal_start:, 1 + j, i]
-                    + mask[horizontal_start:, j]
-                    * (primal_normal_ec[:, :, i] * geofac_div * c_lin_e[c2e, k])[
-                        horizontal_start:, j
-                    ]
-                )
-    return geofac_grg[:, :, 0], geofac_grg[:, :, 1]
+        mask = (e2c[c2e, k] == c2e2c)[horizontal_start:, :]
+        geofac_grg_x[horizontal_start:, 1:] = (
+            geofac_grg_x[horizontal_start:, 1:]
+            + mask * (primal_normal_ec_u * geofac_div * c_lin_e[c2e, k])[horizontal_start:, :]
+        )
+        geofac_grg_y[horizontal_start:, 1:] = (
+            geofac_grg_y[horizontal_start:, 1:]
+            + mask * (primal_normal_ec_v * geofac_div * c_lin_e[c2e, k])[horizontal_start:, :]
+        )
+
+    return geofac_grg_x, geofac_grg_y
 
 
 def compute_geofac_grg(
@@ -273,11 +241,18 @@ def compute_geofac_grg(
     horizontal_start: gtx.int32,
     array_ns: ModuleType = np,
 ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]:
-    primal_normal_ec = functools.partial(_compute_primal_normal_ec, array_ns=array_ns)(
-        primal_normal_cell_x, primal_normal_cell_y, owner_mask, c2e, e2c, horizontal_start
-    )
+    primal_normal_ec_u, primal_normal_ec_v = functools.partial(
+        _compute_primal_normal_ec, array_ns=array_ns
+    )(primal_normal_cell_x, primal_normal_cell_y, owner_mask, c2e, e2c)
     return functools.partial(_compute_geofac_grg, array_ns=array_ns)(
-        primal_normal_ec, geofac_div, c_lin_e, c2e, e2c, c2e2c, horizontal_start
+        primal_normal_ec_u,
+        primal_normal_ec_v,
+        geofac_div,
+        c_lin_e,
+        c2e,
+        e2c,
+        c2e2c,
+        horizontal_start,
     )
 
 
@@ -608,7 +583,7 @@ def _enforce_mass_conservation(
 
     local_summed_weights = array_ns.zeros(c_bln_avg.shape[0])
     residual = array_ns.zeros(c_bln_avg.shape[0])
-    inverse_neighbor_idx = create_inverse_neighbor_index(c2e2c0, array_ns=array_ns)
+    inverse_neighbor_idx = create_inverse_neighbor_index(c2e2c0, c2e2c0, array_ns=array_ns)
 
     for iteration in range(niter):
         local_summed_weights[horizontal_start:] = _compute_local_weights(
@@ -663,13 +638,34 @@ def compute_mass_conserving_bilinear_cell_average_weight(
     )
 
 
-def create_inverse_neighbor_index(c2e2c0, array_ns: ModuleType = np):
-    inv_neighbor_idx = -1 * array_ns.ones(c2e2c0.shape, dtype=int)
+def create_inverse_neighbor_index(source_offset, inverse_offset, array_ns: ModuleType):
+    """
+    The inverse neighbor index determines the position of an central element c_1
+    in the neighbor table of its neighbors:
+
+    For example: for let e_1, e_2, e_3 be the neighboring edges of a cell: c2e(c_1) will
+    map  c_1 -> (e_1, e_2,e_3) then in the inverse lookup table e2c the
+    neighborhoods of e_1, e_2, e_3 will all contain c_1 in some position.
+    Then inverse neighbor index tells what position that is. It essentially says
+    "I am neighbor number x \in (0,1) of my neighboring edges"
+
 
-    for jc in range(c2e2c0.shape[0]):
-        for i in range(c2e2c0.shape[1]):
-            if c2e2c0[jc, i] >= 0:
-                inv_neighbor_idx[jc, i] = array_ns.argwhere(c2e2c0[c2e2c0[jc, i], :] == jc)[0, 0]
+    Args:
+        source_offset:
+        inverse_offset:
+
+    Returns:
+        ndarray of the same shape as target_offset
+
+    """
+    inv_neighbor_idx = -1 * array_ns.ones(inverse_offset.shape, dtype=int)
+
+    for jc in range(inverse_offset.shape[0]):
+        for i in range(inverse_offset.shape[1]):
+            if inverse_offset[jc, i] >= 0:
+                inv_neighbor_idx[jc, i] = array_ns.argwhere(
+                    source_offset[inverse_offset[jc, i], :] == jc
+                )[0, 0]
 
     return inv_neighbor_idx
 

model/common/src/icon4py/model/common/states/prognostic_state.py

@@ -8,9 +8,8 @@
 
 from dataclasses import dataclass
 
-import gt4py.next as gtx
-
-from icon4py.model.common import dimension as dims, field_type_aliases as fa, type_alias as ta
+from icon4py.model.common import field_type_aliases as fa, type_alias as ta
+from icon4py.model.common.dimension import KDim
 
 
 @dataclass
@@ -30,4 +29,4 @@ class PrognosticState:
 
     @property
     def w_1(self) -> fa.CellField[ta.wpfloat]:
-        return gtx.as_field((dims.CellDim,), self.w.ndarray[:, 0])
+        return self.w[KDim(0)]

model/common/tests/interpolation_tests/test_interpolation_fields.py

@@ -148,8 +148,8 @@ def test_compute_geofac_n2s(grid_savepoint, interpolation_savepoint, icon_grid,
 def test_compute_geofac_grg(grid_savepoint, interpolation_savepoint, icon_grid):
     primal_normal_cell_x = grid_savepoint.primal_normal_cell_x().asnumpy()
     primal_normal_cell_y = grid_savepoint.primal_normal_cell_y().asnumpy()
-    geofac_div = interpolation_savepoint.geofac_div()
-    c_lin_e = interpolation_savepoint.c_lin_e()
+    geofac_div = interpolation_savepoint.geofac_div().asnumpy()
+    c_lin_e = interpolation_savepoint.c_lin_e().asnumpy()
     geofac_grg_ref = interpolation_savepoint.geofac_grg()
     owner_mask = grid_savepoint.c_owner_mask()
     c2e = icon_grid.connectivities[dims.C2EDim]
@@ -160,19 +160,25 @@ def test_compute_geofac_grg(grid_savepoint, interpolation_savepoint, icon_grid):
     geofac_grg_0, geofac_grg_1 = compute_geofac_grg(
         primal_normal_cell_x,
         primal_normal_cell_y,
-        owner_mask.asnumpy(),
-        geofac_div.asnumpy(),
-        c_lin_e.asnumpy(),
+        owner_mask,
+        geofac_div,
+        c_lin_e,
         c2e,
         e2c,
         c2e2c,
         horizontal_start,
     )
     assert test_helpers.dallclose(
-        data_alloc.as_numpy(geofac_grg_0), geofac_grg_ref[0].asnumpy(), atol=1e-6, rtol=1e-7
+        data_alloc.as_numpy(geofac_grg_0),
+        geofac_grg_ref[0].asnumpy(),
+        rtol=1e-11,
+        atol=1e-19,
     )
     assert test_helpers.dallclose(
-        data_alloc.as_numpy(geofac_grg_1), geofac_grg_ref[1].asnumpy(), atol=1e-6, rtol=1e-7
+        data_alloc.as_numpy(geofac_grg_1),
+        geofac_grg_ref[1].asnumpy(),
+        rtol=1e-11,
+        atol=1e-19,
     )