Repository: SciTools/iris-esmf-regrid (language: python)

Issue:
regrid_unstructured_to_rectilinear() fails if cube is chunked along a non-horizontal dimension
## 🐛 Bug Report
<!-- Provide a clear description of what the bug is -->

Firstly, apologies for reporting a bug in a non-main branch. I am using the `unstructured_scheme` branch to work with LFRic output.

A `ValueError` occurs if I apply `esmf_regrid.experimental.unstructured_scheme.regrid_unstructured_to_rectilinear()` on cubes that are chunked along the time and/or vertical dimension. This situation is quite common when cubes are loaded from multiple files, so I would really appreciate if you fix this bug or provide a workaround.

My current workaround is to rechunk the data and merge all chunks into one. It works for now, but can become a memory problem for large datasets.

Interestingly, having chunks along the horizontal (mesh) dimension works fine.

Thanks.

## How To Reproduce
Steps to reproduce the behaviour:

```python
import os

import iris
from iris.experimental.ugrid import PARSE_UGRID_ON_LOAD

from esmf_regrid.experimental.unstructured_scheme import (
    regrid_unstructured_to_rectilinear,
)

test_data_dir = iris.config.TEST_DATA_DIR
src_fn = os.path.join(
    test_data_dir, "NetCDF", "unstructured_grid", "lfric_ngvat_3D_1t_half_level_face_grid_derived_theta_in_w3.nc"
)
with PARSE_UGRID_ON_LOAD.context():
    theta = iris.load_cube(src_fn, "air_potential_temperature")

# Rechunk the dask array
theta_chunked = theta.copy(data=theta.core_data().rechunk([1, (19, 19), 864]))

# Load target grid cube.
tgt_fn = os.path.join(
    test_data_dir, "NetCDF", "global", "xyt", "SMALL_hires_wind_u_for_ipcc4.nc"
)
tgt = iris.load_cube(tgt_fn)

# Perform regridding.
result = regrid_unstructured_to_rectilinear(theta_chunked, tgt)
```

<details>
<summary>This results in a `ValueError`: "Dimension 1 has 1 blocks, adjust_chunks specified with 2 blocks"</summary>

```python
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/tmp/ipykernel_23260/3888724239.py in <module>
     16 
     17 # Perform regridding.
---> 18 result = regrid_unstructured_to_rectilinear(theta_chunked, tgt)

~/mambaforge/envs/ideal_exo/lib/python3.8/site-packages/esmf_regrid/experimental/unstructured_scheme.py in regrid_unstructured_to_rectilinear(src_cube, grid_cube, mdtol)
    344     """
    345     regrid_info = _regrid_unstructured_to_rectilinear__prepare(src_cube, grid_cube)
--> 346     result = _regrid_unstructured_to_rectilinear__perform(src_cube, regrid_info, mdtol)
    347     return result
    348 

~/mambaforge/envs/ideal_exo/lib/python3.8/site-packages/esmf_regrid/experimental/unstructured_scheme.py in _regrid_unstructured_to_rectilinear__perform(src_cube, regrid_info, mdtol)
    285     # chunks cover the entire horizontal plane (otherwise they would break
    286     # the regrid function).
--> 287     new_data = _map_complete_blocks(
    288         src_cube,
    289         regrid,

~/mambaforge/envs/ideal_exo/lib/python3.8/site-packages/esmf_regrid/experimental/unstructured_scheme.py in _map_complete_blocks(src, func, dims, out_sizes)
     87         pass
     88 
---> 89     return data.map_blocks(
     90         func, chunks=out_chunks, drop_axis=dropped_dims, dtype=src.dtype
     91     )

~/mambaforge/envs/ideal_exo/lib/python3.8/site-packages/dask/array/core.py in map_blocks(self, func, *args, **kwargs)
   2394     @wraps(map_blocks)
   2395     def map_blocks(self, func, *args, **kwargs):
-> 2396         return map_blocks(func, self, *args, **kwargs)
   2397 
   2398     def map_overlap(self, func, depth, boundary=None, trim=True, **kwargs):

~/mambaforge/envs/ideal_exo/lib/python3.8/site-packages/dask/array/core.py in map_blocks(func, name, token, dtype, chunks, drop_axis, new_axis, meta, *args, **kwargs)
    739         adjust_chunks = None
    740 
--> 741     out = blockwise(
    742         func,
    743         out_ind,

~/mambaforge/envs/ideal_exo/lib/python3.8/site-packages/dask/array/blockwise.py in blockwise(func, out_ind, name, token, dtype, adjust_chunks, new_axes, align_arrays, concatenate, meta, *args, **kwargs)
    260                 elif isinstance(adjust_chunks[ind], (tuple, list)):
    261                     if len(adjust_chunks[ind]) != len(chunks[i]):
--> 262                         raise ValueError(
    263                             f"Dimension {i} has {len(chunks[i])} blocks, adjust_chunks "
    264                             f"specified with {len(adjust_chunks[ind])} blocks"

ValueError: Dimension 1 has 1 blocks, adjust_chunks specified with 2 blocks
```

</details>


## Expected behaviour
<!-- A clear and concise description of what you expected to happen -->
Regridding that works on an arbitrarily chunked array.

## Environment 
 - OS and version: Ubuntu 20.04 LTS
 - `esmf_regrid` version: 0.3.dev0
 - `iris` version: 3.2.dev0



Discussion / hints:
# ✨ Congratulations! ✨ Thanks for submitting your first issue to `iris-esmf-regrid`. We really appreciate it and will get back to you as soon as possible. Awesome job 👍

Fix patch:
diff --git a/esmf_regrid/experimental/unstructured_scheme.py b/esmf_regrid/experimental/unstructured_scheme.py
index e197f613..4a2bd6ab 100644
--- a/esmf_regrid/experimental/unstructured_scheme.py
+++ b/esmf_regrid/experimental/unstructured_scheme.py
@@ -63,22 +63,34 @@ def _map_complete_blocks(src, func, dims, out_sizes):
     num_dims = len(dims)
     num_out = len(out_sizes)
     dropped_dims = []
+    new_axis = None
     if num_out > num_dims:
         # While this code should be robust for cases where num_out > num_dims > 1,
         # there is some ambiguity as to what their behaviour ought to be.
         # Since these cases are out of our own scope, we explicitly ignore them
         # for the time being.
         assert num_dims == 1
+        # While this code should be robust for cases where num_out > 2,
+        # we expect to handle at most 2D grids.
+        # Since these cases are out of our own scope, we explicitly ignore them
+        # for the time being.
+        assert num_out == 2
         slice_index = sorted_dims[-1]
         # Insert the remaining contents of out_sizes in the position immediately
         # after the last dimension.
         out_chunks[slice_index:slice_index] = out_sizes[num_dims:]
+        new_axis = range(slice_index, slice_index + num_out - num_dims)
     elif num_dims > num_out:
         # While this code should be robust for cases where num_dims > num_out > 1,
         # there is some ambiguity as to what their behaviour ought to be.
         # Since these cases are out of our own scope, we explicitly ignore them
         # for the time being.
         assert num_out == 1
+        # While this code should be robust for cases where num_dims > 2,
+        # we expect to handle at most 2D grids.
+        # Since these cases are out of our own scope, we explicitly ignore them
+        # for the time being.
+        assert num_dims == 2
         dropped_dims = sorted_dims[num_out:]
         # Remove the remaining dimensions from the expected output shape.
         for dim in dropped_dims[::-1]:
@@ -87,7 +99,11 @@ def _map_complete_blocks(src, func, dims, out_sizes):
         pass
 
     return data.map_blocks(
-        func, chunks=out_chunks, drop_axis=dropped_dims, dtype=src.dtype
+        func,
+        chunks=out_chunks,
+        drop_axis=dropped_dims,
+        new_axis=new_axis,
+        dtype=src.dtype,
     )
 
 


Test patch:
diff --git a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py
index 9546ac23..d35e79e2 100644
--- a/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py
+++ b/esmf_regrid/tests/unit/experimental/unstructured_scheme/test_MeshToGridESMFRegridder.py
@@ -141,17 +141,20 @@ def test_laziness():
     n_lons = 12
     n_lats = 10
     h = 4
+    i = 9
     lon_bounds = (-180, 180)
     lat_bounds = (-90, 90)
 
     mesh = _gridlike_mesh(n_lons, n_lats)
 
-    src_data = np.arange(n_lats * n_lons * h).reshape([-1, h])
-    src_data = da.from_array(src_data, chunks=[15, 2])
+    # Add a chunked dimension both before and after the mesh dimension.
+    # The leading length 1 dimension matches the example in issue #135.
+    src_data = np.arange(i * n_lats * n_lons * h).reshape([1, i, -1, h])
+    src_data = da.from_array(src_data, chunks=[1, 3, 15, 2])
     src = Cube(src_data)
     mesh_coord_x, mesh_coord_y = mesh.to_MeshCoords("face")
-    src.add_aux_coord(mesh_coord_x, 0)
-    src.add_aux_coord(mesh_coord_y, 0)
+    src.add_aux_coord(mesh_coord_x, 2)
+    src.add_aux_coord(mesh_coord_y, 2)
     tgt = _grid_cube(n_lons, n_lats, lon_bounds, lat_bounds, circular=True)
 
     rg = MeshToGridESMFRegridder(src, tgt)
@@ -160,6 +163,6 @@ def test_laziness():
     result = rg(src)
     assert result.has_lazy_data()
     out_chunks = result.lazy_data().chunks
-    expected_chunks = ((10,), (12,), (2, 2))
+    expected_chunks = ((1,), (3, 3, 3), (10,), (12,), (2, 2))
     assert out_chunks == expected_chunks
-    assert np.allclose(result.data.reshape([-1, h]), src_data)
+    assert np.allclose(result.data.reshape([1, i, -1, h]), src_data)