Add single-rank lsq pseudoinv factory test (#1099)

Related to
#1065 (comment). This
uses @nfarabullini's fixes for doing the LSQ_PSEUDOINV computation with
a GPU backend from #1012 to have them merged separately.

This also adds LSQ_PSEUDOINV to the parallel test in
test_parallel_grid_manager.py. This required changing the computation to
avoid doing the SVD on halo points.

Custom dimensions are added for LSQ_PSEUDOINV to correctly declare the
dimensions in the factory. Previously the dimension was equivalne to a
scalar, which lead to the factory returning the plain numpy array. With
this change the mypy type ignores can be removed from
test_parallel_interpolation.py.

---------

Co-authored-by: Nicoletta Farabullini <41536517+nfarabullini@users.noreply.github.com>

Author: msimberg

model/common/src/icon4py/model/common/dimension.py

@@ -16,6 +16,8 @@
 CellDim = gtx.Dimension("Cell")
 VertexDim = gtx.Dimension("Vertex")
 MAIN_HORIZONTAL_DIMENSIONS = {"CellDim": CellDim, "EdgeDim": EdgeDim, "VertexDim": VertexDim}
+LsqCDim = gtx.Dimension("LsqC", gtx.DimensionKind.LOCAL)
+LsqUnkDim = gtx.Dimension("LsqUnk", gtx.DimensionKind.LOCAL)
 E2CDim = gtx.Dimension("E2C", gtx.DimensionKind.LOCAL)
 E2VDim = gtx.Dimension("E2V", gtx.DimensionKind.LOCAL)
 C2EDim = gtx.Dimension("C2E", gtx.DimensionKind.LOCAL)

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

@@ -243,7 +243,7 @@ def _register_computed_fields(self) -> None:
                 array_ns=self._xp,
             ),
             fields=(attrs.LSQ_PSEUDOINV,),
-            domain=(),
+            domain=(dims.CellDim, dims.LsqUnkDim, dims.LsqCDim),
             deps={
                 "cell_center_x": geometry_attrs.CELL_CENTER_X,
                 "cell_center_y": geometry_attrs.CELL_CENTER_Y,

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

@@ -1167,8 +1167,8 @@ def compute_lsq_pseudoinv(
     for jjb in range(lsq_dim_c):
         for jjk in range(lsq_dim_unk):
             for jc in range(start_idx, min_rlcell_int):
-                u, s, v_t = array_ns.linalg.svd(z_lsq_mat_c[jc, :, :])
                 if cell_owner_mask[jc]:
+                    u, s, v_t = array_ns.linalg.svd(z_lsq_mat_c[jc, :, :])
                     lsq_pseudoinv[jc, :lsq_dim_unk, jjb] = (
                         lsq_pseudoinv[jc, :lsq_dim_unk, jjb]
                         + v_t[jjk, :lsq_dim_unk] / s[jjk] * u[jjb, jjk] * lsq_weights_c[jc, jjb]
@@ -1237,7 +1237,11 @@ def compute_lsq_coeffs(
             for js in range(lsq_dim_stencil):
                 z_dist_g[:, js, :] = array_ns.asarray(
                     gnomonic_proj(
-                        cell_lon, cell_lat, cell_lon[c2e2c[:, js]], cell_lat[c2e2c[:, js]]
+                        cell_lon,
+                        cell_lat,
+                        cell_lon[c2e2c[:, js]],
+                        cell_lat[c2e2c[:, js]],
+                        array_ns,
                     )
                 ).T
 
@@ -1252,15 +1256,17 @@ def compute_lsq_coeffs(
                 ilc_s = c2e2c[jc, :lsq_dim_stencil]
                 cc_cell = array_ns.zeros((lsq_dim_stencil, 2))
 
-                cc_cv = (cell_center_x[jc], cell_center_y[jc])
+                cc_cv = array_ns.asarray((cell_center_x[jc], cell_center_y[jc]))
                 for js in range(lsq_dim_stencil):
-                    cc_cell[js, :] = diff_on_edges_torus_numpy(
-                        cell_center_x[jc],
-                        cell_center_y[jc],
-                        cell_center_x[ilc_s][js],
-                        cell_center_y[ilc_s][js],
-                        domain_length,
-                        domain_height,
+                    cc_cell[js, :] = array_ns.asarray(
+                        diff_on_edges_torus_numpy(
+                            cell_center_x[jc],
+                            cell_center_y[jc],
+                            cell_center_x[ilc_s][js],
+                            cell_center_y[ilc_s][js],
+                            domain_length,
+                            domain_height,
+                        )
                     )
                 z_dist_g[jc, :, :] = cc_cell - cc_cv
 

model/common/src/icon4py/model/common/math/projection.py

@@ -5,7 +5,7 @@
 #
 # Please, refer to the LICENSE file in the root directory.
 # SPDX-License-Identifier: BSD-3-Clause
-
+from types import ModuleType
 
 import numpy as np
 
@@ -17,6 +17,7 @@ def gnomonic_proj(
     lat_c: data_alloc.NDArray,
     lon: data_alloc.NDArray,
     lat: data_alloc.NDArray,
+    array_ns: ModuleType = np,
 ) -> tuple[data_alloc.NDArray, data_alloc.NDArray]:
     """
     Compute gnomonic projection.
@@ -38,11 +39,16 @@ def gnomonic_proj(
     TODO:
         replace this with a suitable library call
     """
-    cosc = np.sin(lat_c) * np.sin(lat) + np.cos(lat_c) * np.cos(lat) * np.cos(lon - lon_c)
+    cosc = array_ns.sin(lat_c) * array_ns.sin(lat) + array_ns.cos(lat_c) * array_ns.cos(
+        lat
+    ) * array_ns.cos(lon - lon_c)
     zk = 1.0 / cosc
 
-    x = zk * np.cos(lat) * np.sin(lon - lon_c)
-    y = zk * (np.cos(lat_c) * np.sin(lat) - np.sin(lat_c) * np.cos(lat) * np.cos(lon - lon_c))
+    x = zk * array_ns.cos(lat) * array_ns.sin(lon - lon_c)
+    y = zk * (
+        array_ns.cos(lat_c) * array_ns.sin(lat)
+        - array_ns.sin(lat_c) * array_ns.cos(lat) * array_ns.cos(lon - lon_c)
+    )
 
     return x, y
 

model/common/tests/common/grid/mpi_tests/test_parallel_grid_manager.py

@@ -88,6 +88,9 @@ def _get_neighbor_tables(grid: base.Grid) -> dict:
 def gather_field(field: np.ndarray, props: decomp_defs.ProcessProperties) -> tuple:
     constant_dims = tuple(field.shape[1:])
     _log.info(f"gather_field on rank={props.rank} - gathering field of local shape {field.shape}")
+    # Because of sparse indexing the field may have a non-contigous layout,
+    # which Gatherv doesn't support. Make sure the field is contiguous.
+    field = np.ascontiguousarray(field)
     constant_length = functools.reduce(operator.mul, constant_dims, 1)
     local_sizes = np.array(props.comm.gather(field.size, root=0))
     if props.rank == 0:
@@ -337,6 +340,7 @@ def test_geometry_fields_compare_single_multi_rank(
         interpolation_attributes.GEOFAC_GRG_Y,
         interpolation_attributes.GEOFAC_N2S,
         interpolation_attributes.GEOFAC_ROT,
+        interpolation_attributes.LSQ_PSEUDOINV,
         interpolation_attributes.NUDGECOEFFS_E,
         interpolation_attributes.POS_ON_TPLANE_E_X,
         interpolation_attributes.POS_ON_TPLANE_E_Y,

model/common/tests/common/grid/unit_tests/test_icon.py

@@ -214,6 +214,8 @@ def test_when_replace_skip_values_then_only_pentagon_points_remain(
 ) -> None:
     if dim == dims.V2E2VDim:
         pytest.skip("V2E2VDim is not supported in the current grid configuration.")
+    if dim in (dims.LsqCDim, dims.LsqUnkDim):
+        pytest.skip("LsqCDim and LsqUnkDim are not offset dimensions.")
     grid = utils.run_grid_manager(grid_descriptor, keep_skip_values=False, backend=backend).grid
     connectivity = grid.get_connectivity(dim.value)
     if dim in icon.CONNECTIVITIES_ON_PENTAGONS and not grid.limited_area:

model/common/tests/common/interpolation/mpi_tests/test_parallel_interpolation.py

@@ -235,9 +235,8 @@ def test_distributed_interpolation_lsq_pseudoinv(
     parallel_helpers.log_process_properties(processor_props)
     parallel_helpers.log_local_field_size(decomposition_info)
     factory = interpolation_factory_from_savepoint
-    field_ref_1 = interpolation_savepoint.__getattribute__("lsq_pseudoinv_1")().asnumpy()
-    field_ref_2 = interpolation_savepoint.__getattribute__("lsq_pseudoinv_2")().asnumpy()
-    field_1 = factory.get(attrs.LSQ_PSEUDOINV)[:, 0, :]
-    field_2 = factory.get(attrs.LSQ_PSEUDOINV)[:, 1, :]
-    assert test_utils.dallclose(field_1, field_ref_1, atol=1e-15)  # type: ignore[arg-type] # mypy does not recognize sliced array as still an array
-    assert test_utils.dallclose(field_2, field_ref_2, atol=1e-15)  # type: ignore[arg-type] # mypy does not recognize sliced array as still an array
+    field_ref_1 = interpolation_savepoint.lsq_pseudoinv_1().asnumpy()
+    field_ref_2 = interpolation_savepoint.lsq_pseudoinv_2().asnumpy()
+    field = factory.get(attrs.LSQ_PSEUDOINV).asnumpy()
+    assert test_utils.dallclose(field[:, 0, :], field_ref_1, atol=1e-15)
+    assert test_utils.dallclose(field[:, 1, :], field_ref_2, atol=1e-15)

model/common/tests/common/interpolation/unit_tests/test_interpolation_factory.py

@@ -374,3 +374,18 @@ def test_rbf_interpolation_coeffs_vertex(
         field_v2[horizontal_start:],
         atol=RBF_TOLERANCES[dims.VertexDim][experiment.name],
     )
+
+
+@pytest.mark.level("integration")
+@pytest.mark.datatest
+def test_lsq_pseudoinv(
+    interpolation_savepoint: serialbox.InterpolationSavepoint,
+    experiment: definitions.Experiment,
+    backend: gtx_typing.Backend | None,
+) -> None:
+    field_ref_1 = interpolation_savepoint.lsq_pseudoinv_1().asnumpy()
+    field_ref_2 = interpolation_savepoint.lsq_pseudoinv_2().asnumpy()
+    factory = _get_interpolation_factory(backend, experiment)
+    field = factory.get(attrs.LSQ_PSEUDOINV).asnumpy()
+    assert test_helpers.dallclose(field_ref_1, field[:, 0, :], atol=1e-15)
+    assert test_helpers.dallclose(field_ref_2, field[:, 1, :], atol=1e-15)

model/testing/src/icon4py/model/testing/serialbox.py

@@ -703,10 +703,10 @@ def rbf_vec_idx_v(self):
         return self._get_field("rbf_vec_idx_v", dims.VertexDim, dims.V2EDim)
 
     def lsq_pseudoinv_1(self):
-        return self._get_field("lsq_pseudoinv_1", dims.CellDim, dims.C2E2CDim)
+        return self._get_field("lsq_pseudoinv_1", dims.CellDim, dims.LsqCDim)
 
     def lsq_pseudoinv_2(self):
-        return self._get_field("lsq_pseudoinv_2", dims.CellDim, dims.C2E2CDim)
+        return self._get_field("lsq_pseudoinv_2", dims.CellDim, dims.LsqCDim)
 
 
 class MetricSavepoint(IconSavepoint):