WIP: Greenline metrics coefficients aj

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

@@ -0,0 +1,124 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+from gt4py.next import where
+from gt4py.next.ffront.decorator import field_operator
+from gt4py.next.ffront.fbuiltins import Field
+
+from icon4py.model.common.dimension import C2E, V2E, C2EDim, CellDim, EdgeDim, V2EDim, VertexDim
+
+def grad_fd_norm(
+    psi_c: np.array,
+    inv_dual_edge_length: np.array,
+    e2c: np.array,
+    second_boundary_layer_start_index: np.int32,
+    nlev,
+) -> np.array:
+    llb = second_boundary_layer_start_index
+    grad_norm_psi_e = np.zeros([len(e2c[:, 0]), nlev])
+    for i in range(nlev):
+        grad_norm_psi_e[llb:, i] = (psi_c[e2c[llb:, 1], i] - psi_c[e2c[llb:, 0], i]) * inv_dual_edge_length[llb:]
+    return grad_norm_psi_e
+
+def grad_fd_tang(
+    psi_v: np.array,
+    inv_primal_edge_length: np.array,
+    tangent_orientation: np.array,
+    e2v: np.array,
+    third_boundary_layer_start_index: np.int32,
+    nlev,
+) -> np.array:
+    llb = third_boundary_layer_start_index
+    grad_tang_psi_e = np.zeros([e2v.shape[0], nlev])
+    for i in range(nlev):
+        grad_tang_psi_e[llb:, i] = tangent_orientation[llb:] * (psi_v[e2v[llb:, 1], i] - psi_v[e2v[llb:, 0], i]) * inv_primal_edge_length[llb:]
+    return grad_tang_psi_e
+
+def compute_ddxn_z_half_e(
+    z_ifc: np.array,
+    inv_dual_edge_length: np.array,
+    e2c: np.array,
+    second_boundary_layer_start_index: np.int32,
+) -> np.array:
+    nlev = z_ifc.shape[1]
+    ddxn_z_half_e = grad_fd_norm(z_ifc, inv_dual_edge_length, e2c, second_boundary_layer_start_index, nlev)
+    return ddxn_z_half_e
+
+def compute_ddxt_z_half_e(
+    z_ifv: np.array,
+    inv_primal_edge_length: np.array,
+    tangent_orientation: np.array,
+    e2v: np.array,
+    third_boundary_layer_start_index: np.int32,
+) -> np.array:
+    nlev = z_ifv.shape[1]
+    ddxt_z_half_e = grad_fd_tang(z_ifv, inv_primal_edge_length, tangent_orientation, e2v, third_boundary_layer_start_index, nlev)
+    return ddxt_z_half_e
+
+def compute_ddxnt_z_full(
+    z_ddxnt_z_half_e: np.array,
+) -> np.array:
+    ddxnt_z_full = 0.5 * (z_ddxnt_z_half_e[:, :z_ddxnt_z_half_e.shape[1]-1] + z_ddxnt_z_half_e[:, 1:])
+    return ddxnt_z_full
+
+def compute_cells2verts_scalar(
+    p_cell_in: np.array,
+    c_int: np.array,
+    v2c: np.array,
+    second_boundary_layer_start_index: np.int32,
+) -> np.array:
+    llb = second_boundary_layer_start_index
+    kdim = p_cell_in.shape[1]
+    p_vert_out = np.zeros([c_int.shape[0], kdim])
+    for j in range(kdim):
+        for i in range(6):
+            p_vert_out[llb:, j] = p_vert_out[llb:, j] + c_int[llb:, i] * p_cell_in[v2c[llb:, i], j]
+    return p_vert_out
+
+def compute_cells_aw_verts(
+    dual_area: np.array,
+    edge_vert_length: np.array,
+    edge_cell_length: np.array,
+    owner_mask: np.array,
+    e2c: np.array,
+    v2c: np.array,
+    v2e: np.array,
+    e2v: np.array,
+    second_boundary_layer_start_index: np.int32,
+    second_boundary_layer_end_index: np.int32,
+) -> np.array:
+    llb = second_boundary_layer_start_index
+    cells_aw_verts = np.zeros([second_boundary_layer_end_index, 6])
+    index = np.zeros([second_boundary_layer_end_index, 6])
+    for j in range (6):
+        index[:, j] = np.arange(second_boundary_layer_end_index)
+    idx_ve = np.where(index == e2v[v2e, 0], 0, 1)
+
+    for i in range(2):
+        for j in range (6):
+            for k in range (6):
+                cells_aw_verts[llb:, k] = np.where(np.logical_and(owner_mask[:], e2c[v2e[llb:, j], i] == v2c[llb:, k]), cells_aw_verts[llb:, k] + 0.5 / dual_area[llb:] * edge_vert_length[v2e[llb:, j], idx_ve[llb:, j]] * edge_cell_length[v2e[llb:, j], i], cells_aw_verts[llb:, k])
+    return cells_aw_verts

model/common/src/icon4py/model/common/test_utils/serialbox_utils.py

@@ -134,6 +134,12 @@ def _read(self, name: str, offset=0, dtype=int):
 
 
 class IconGridSavepoint(IconSavepoint):
+    def v_dual_area(self):
+        return self._get_field("v_dual_area", VertexDim)
+
+    def edge_vert_length(self):
+        return self._get_field("edge_vert_length", EdgeDim, E2C2VDim)
+
     def vct_a(self):
         return self._get_field("vct_a", KDim)
 
@@ -211,6 +217,9 @@ def edge_end_index(self):
         # one off accounts for being exclusive [from:to)
         return self.serializer.read("e_end_index", self.savepoint)
 
+    def v_owner_mask(self):
+        return self._get_field("v_owner_mask", VertexDim, dtype=bool)
+
     def c_owner_mask(self):
         return self._get_field("c_owner_mask", CellDim, dtype=bool)
 

model/common/tests/interpolation_tests/test_interpolation_fields3.py

@@ -0,0 +1,152 @@
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+# ICON4Py - ICON inspired code in Python and GT4Py
+#
+# Copyright (c) 2022, ETH Zurich and MeteoSwiss
+# All rights reserved.
+#
+# This file is free software: you can redistribute it and/or modify it under
+# the terms of the GNU General Public License as published by the
+# Free Software Foundation, either version 3 of the License, or any later
+# version. See the LICENSE.txt file at the top-level directory of this
+# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import numpy as np
+import pytest
+from gt4py.next.iterator.builtins import int32
+
+from icon4py.model.common.dimension import (
+    C2E2CDim,
+    C2EDim,
+    CellDim,
+    E2C2EDim,
+    E2CDim,
+    E2VDim,
+    EdgeDim,
+    V2CDim,
+    V2EDim,
+    VertexDim,
+)
+from icon4py.model.common.grid.horizontal import HorizontalMarkerIndex
+from icon4py.model.common.interpolation.interpolation_fields3 import (
+    compute_ddxn_z_half_e,
+    compute_ddxnt_z_full,
+    compute_cells_aw_verts,
+    compute_cells2verts_scalar,
+    compute_ddxt_z_half_e,
+)
+from icon4py.model.common.test_utils.datatest_fixtures import (  # noqa: F401  # import fixtures from test_utils package
+    data_provider,
+    datapath,
+    download_ser_data,
+    experiment,
+    processor_props,
+    ranked_data_path,
+)
+from icon4py.model.common.test_utils.helpers import zero_field
+
+
+@pytest.mark.datatest
+def test_compute_ddxn_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint):  # fixture
+    z_ifc = metrics_savepoint.z_ifc().asnumpy()
+    inv_dual_edge_length = grid_savepoint.inv_dual_edge_length().asnumpy()
+    e2c = icon_grid.connectivities[E2CDim]
+#    edge_cell_length = grid_savepoint.edge_cell_length()
+#    owner_mask = grid_savepoint.e_owner_mask()
+    ddxn_z_full_ref = metrics_savepoint.ddxn_z_full().asnumpy()
+    lateral_boundary = np.arange(2)
+    lateral_boundary[0] = icon_grid.get_start_index(
+        EdgeDim,
+        HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 1,
+    )
+    lateral_boundary[1] = icon_grid.get_end_index(
+        EdgeDim,
+        HorizontalMarkerIndex.lateral_boundary(EdgeDim) - 1,
+    )
+    ddxn_z_half_e = compute_ddxn_z_half_e(
+        z_ifc,
+        inv_dual_edge_length,
+        e2c,
+        lateral_boundary[0],
+    )
+    ddxn_z_full = compute_ddxnt_z_full(
+        ddxn_z_half_e,
+    )
+
+    assert np.allclose(ddxn_z_full, ddxn_z_full_ref)
+
+
+@pytest.mark.datatest
+def test_compute_ddxt_z_full_e(grid_savepoint, interpolation_savepoint, icon_grid, metrics_savepoint):
+    z_ifc = metrics_savepoint.z_ifc().asnumpy()
+    dual_area = grid_savepoint.v_dual_area().asnumpy()
+    edge_vert_length = grid_savepoint.edge_vert_length().asnumpy()
+    edge_cell_length = grid_savepoint.edge_cell_length().asnumpy()
+    tangent_orientation = grid_savepoint.tangent_orientation().asnumpy()
+    inv_primal_edge_length = grid_savepoint.inverse_primal_edge_lengths().asnumpy()
+    owner_mask = grid_savepoint.v_owner_mask()
+    ddxt_z_full_ref = metrics_savepoint.ddxt_z_full().asnumpy()
+    e2c = icon_grid.connectivities[E2CDim]
+    v2c = icon_grid.connectivities[V2CDim]
+    v2e = icon_grid.connectivities[V2EDim]
+    e2v = icon_grid.connectivities[E2VDim]
+    second_boundary_layer_start_index_vertex = icon_grid.get_start_index(
+        VertexDim,
+        HorizontalMarkerIndex.lateral_boundary(VertexDim) + 1,
+    )
+    second_boundary_layer_end_index_vertex = icon_grid.get_end_index(
+        VertexDim,
+        HorizontalMarkerIndex.lateral_boundary(VertexDim) - 1,
+    )
+    second_boundary_layer_start_index_cell = icon_grid.get_start_index(
+        CellDim,
+        HorizontalMarkerIndex.lateral_boundary(CellDim) + 1,
+    )
+    second_boundary_layer_end_index_cell = icon_grid.get_end_index(
+        CellDim,
+        HorizontalMarkerIndex.lateral_boundary(CellDim) - 1,
+    )
+    third_boundary_layer_start_index_edge = icon_grid.get_start_index(
+        EdgeDim,
+        HorizontalMarkerIndex.lateral_boundary(EdgeDim) + 2,
+    )
+    cells_aw_verts = compute_cells_aw_verts(
+        dual_area,
+        edge_vert_length,
+        edge_cell_length,
+        owner_mask,
+        e2c,
+        v2c,
+        v2e,
+        e2v,
+        second_boundary_layer_start_index_vertex,
+        second_boundary_layer_end_index_vertex,
+    )
+    cells_aw_verts_ref = interpolation_savepoint.c_intp().asnumpy()
+    assert np.allclose(cells_aw_verts, cells_aw_verts_ref)
+
+    z_ifv = compute_cells2verts_scalar(z_ifc, cells_aw_verts, v2c, second_boundary_layer_start_index_vertex)
+    ddxt_z_half_e = compute_ddxt_z_half_e(
+        z_ifv,
+        inv_primal_edge_length,
+        tangent_orientation,
+        e2v,
+        third_boundary_layer_start_index_edge,
+    )
+    ddxt_z_full = compute_ddxnt_z_full(
+        ddxt_z_half_e,
+    )
+
+    assert np.allclose(ddxt_z_full, ddxt_z_full_ref)