Fix curvilinear lazy regridding (#242)

* fix curvilinear lazy regridding

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review comments

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

esmf_regrid/experimental/unstructured_scheme.py

@@ -220,7 +220,8 @@ def _regrid_unstructured_to_rectilinear__perform(src_cube, regrid_info, mdtol):
     if len(grid_x.shape) == 1:
         chunk_shape = (len(grid_x.points), len(grid_y.points))
     else:
-        chunk_shape = grid_x.shape
+        # Due to structural reasons, the order here must be reversed.
+        chunk_shape = grid_x.shape[::-1]
     new_data = _map_complete_blocks(
         src_cube,
         regrid,

esmf_regrid/schemes.py

@@ -235,7 +235,8 @@ def _regrid_rectilinear_to_rectilinear__perform(src_cube, regrid_info, mdtol):
     if len(grid_x.shape) == 1:
         chunk_shape = (len(grid_x.points), len(grid_y.points))
     else:
-        chunk_shape = grid_x.shape
+        # Due to structural reasons, the order here must be reversed.
+        chunk_shape = grid_x.shape[::-1]
     new_data = map_complete_blocks(
         src_cube,
         regrid,

esmf_regrid/tests/unit/experimental/unstructured_scheme/test_GridToMeshESMFRegridder.py

@@ -404,7 +404,9 @@ def test_curvilinear():
     cube.add_dim_coord(height, 3)
 
     regridder = GridToMeshESMFRegridder(grid, tgt)
+    cube_lazy = cube.copy(da.array(cube.data))
     result = regridder(cube)
+    result_lazy = regridder(cube_lazy)
 
     # Lenient check for data.
     expected_data = np.empty([t, mesh_length, h])
@@ -421,3 +423,4 @@ def test_curvilinear():
     # Check metadata and scalar coords.
     result.data = expected_data
     assert expected_cube == result
+    assert result_lazy == result

esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py

@@ -382,6 +382,8 @@ def test_curvilinear():
     src_cube.transpose([1, 0, 2])
     regridder = MeshToGridESMFRegridder(src_cube, tgt)
     result = regridder(mesh_cube)
+    mesh_cube_lazy = mesh_cube.copy(da.array(mesh_cube.data))
+    result_lazy = regridder(mesh_cube_lazy)
 
     # Lenient check for data.
     expected_data = np.empty([t, n_lats, n_lons, h])
@@ -397,3 +399,4 @@ def test_curvilinear():
     # Check metadata and scalar coords.
     result.data = expected_data
     assert expected_cube == result
+    assert result_lazy == result

esmf_regrid/tests/unit/schemes/test_regrid_rectilinear_to_rectilinear.py

@@ -177,6 +177,65 @@ def test_laziness():
     assert np.allclose(result.data, src_data)
 
 
+def test_laziness_curvilinear():
+    """
+    Test for :func:`esmf_regrid.schemes.regrid_rectilinear_to_rectilinear`.
+
+    Tests the handling of lazy data. Ensures that handling is the same
+    for non-lazy data.
+    """
+    h = 2
+    t = 4
+    e = 6
+    src_lats = 3
+    src_lons = 5
+
+    tgt_lats = 5
+    tgt_lons = 3
+
+    lon_bounds = (-180, 180)
+    lat_bounds = (-90, 90)
+
+    src_grid = _curvilinear_cube(
+        src_lons,
+        src_lats,
+        lon_bounds,
+        lat_bounds,
+    )
+    tgt_grid = _curvilinear_cube(
+        tgt_lons,
+        tgt_lats,
+        lon_bounds,
+        lat_bounds,
+    )
+
+    height = DimCoord(np.arange(h), standard_name="height")
+    time = DimCoord(np.arange(t), standard_name="time")
+    extra = AuxCoord(np.arange(e), long_name="extra dim")
+
+    src_data = np.empty([h, src_lats, t, src_lons, e])
+    src_data[:] = np.arange(t * h * e).reshape([h, t, e])[
+        :, np.newaxis, :, np.newaxis, :
+    ]
+    src_data_lazy = da.array(src_data)
+
+    src_cube = Cube(src_data)
+    src_cube.add_dim_coord(height, 0)
+    src_cube.add_aux_coord(src_grid.coord("latitude"), (1, 3))
+    src_cube.add_dim_coord(time, 2)
+    src_cube.add_aux_coord(src_grid.coord("longitude"), (1, 3))
+    src_cube.add_aux_coord(extra, 4)
+
+    src_cube_lazy = src_cube.copy(src_data_lazy)
+
+    result = regrid_rectilinear_to_rectilinear(src_cube, tgt_grid)
+    result_lazy = regrid_rectilinear_to_rectilinear(src_cube_lazy, tgt_grid)
+
+    assert result_lazy.has_lazy_data()
+
+    assert result_lazy == result
+
+
 def test_extra_dims_curvilinear():
     """
     Test for :func:`esmf_regrid.schemes.regrid_rectilinear_to_rectilinear`.