regrid_unstructured_to_rectilinear() fails if cube is chunked along a non-horizontal dimension

[dennissergeev]

🐛 Bug Report

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:

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)

This results in a `ValueError`: "Dimension 1 has 1 blocks, adjust_chunks specified with 2 blocks"

```python --------------------------------------------------------------------------- ValueError Traceback (most recent call last) /tmp/ipykernel_23260/3888724239.py in 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 ```

Expected behaviour

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

[github-actions[bot]]

✨ 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 👍

[stephenworsley]

Thanks for highlighting this @dennissergeev , I'm working on a fix in #136.

[dennissergeev]

thanks very much for the quick fix @stephenworsley !