Merge branch 'main' into fix_for_esmpy_8.6

* main:
  Fix for esmpy (SciTools#353)
  Fix usage of map_blocks (SciTools#338)

# Conflicts:
#	.github/workflows/ci-tests.yml
#	.github/workflows/ci-wheels.yml
#	esmf_regrid/experimental/unstructured_regrid.py
#	noxfile.py
#	pyproject.toml
#	requirements/py312.yml

.github/workflows/ci-tests.yml

@@ -94,11 +94,12 @@ jobs:
       - name: "tests (py${{ matrix.python-version }})"
         env:
           PY_VER: ${{ matrix.python-version }}
+          COVERAGE: ${{ matrix.coverage }}
         run: |
           nox --session tests -- --verbose
 
       - name: Upload coverage report
         uses: codecov/codecov-action@v4
+        if: ${{ matrix.coverage }}
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
-        if: ${{ matrix.coverage }}

esmf_regrid/experimental/unstructured_regrid.py

@@ -95,10 +95,10 @@ def _as_esmf_info(self):
         # the data must be translated into a form ESMF understands
         num_node = self.node_coords.shape[0]
         num_elem = self.fnc.shape[0]
-        nodeId = np.array(range(num_node)) + 1
+        nodeId = np.arange(1, num_node + 1)
         nodeCoord = self.node_coords.flatten()
         nodeOwner = np.zeros([num_node])  # regridding currently serial
-        elemId = np.array(range(num_elem)) + 1
+        elemId = np.arange(1, num_elem + 1)
         elemType = self.fnc.count(axis=1)
         # Experiments seem to indicate that ESMF is using 0 indexing here
         elemConn = self.fnc.compressed() - self.nsi

esmf_regrid/schemes.py

@@ -2,7 +2,6 @@
 
 from collections import namedtuple
 import copy
-import functools
 
 from cf_units import Unit
 import dask.array as da
@@ -239,16 +238,16 @@ def _mesh_to_MeshInfo(mesh, location, mask=None):
     return meshinfo
 
 
-def _regrid_along_dims(regridder, data, dims, num_out_dims, mdtol):
+def _regrid_along_dims(data, regridder, dims, num_out_dims, mdtol):
     """
     Perform regridding on data over specific dimensions.
 
     Parameters
     ----------
-    regridder : Regridder
-        An instance of Regridder initialised to perform regridding.
     data : array
         The data to be regrid.
+    regridder : Regridder
+        An instance of Regridder initialised to perform regridding.
     dims : tuple of int
         The dimensions in the source data to regrid along.
     num_out_dims : int
@@ -276,7 +275,7 @@ def _regrid_along_dims(regridder, data, dims, num_out_dims, mdtol):
     return result
 
 
-def _map_complete_blocks(src, func, dims, out_sizes):
+def _map_complete_blocks(src, func, active_dims, out_sizes, *args, **kwargs):
     """
     Apply a function to complete blocks.
 
@@ -296,7 +295,7 @@ def _map_complete_blocks(src, func, dims, out_sizes):
         Source :class:`~iris.cube.Cube` that function is applied to.
     func : function
         Function to apply.
-    dims : tuple of int
+    active_dims : tuple of int
         Dimensions that cannot be chunked.
     out_sizes : tuple of int
         Output size of dimensions that cannot be chunked.
@@ -309,25 +308,25 @@ def _map_complete_blocks(src, func, dims, out_sizes):
 
     """
     if not src.has_lazy_data():
-        return func(src.data)
+        return func(src.data, *args, **kwargs)
 
     data = src.lazy_data()
 
     # Ensure dims are not chunked
     in_chunks = list(data.chunks)
-    for dim in dims:
+    for dim in active_dims:
         in_chunks[dim] = src.shape[dim]
     data = data.rechunk(in_chunks)
 
     # Determine output chunks
-    num_dims = len(dims)
+    num_dims = len(active_dims)
     num_out = len(out_sizes)
     out_chunks = list(data.chunks)
-    sorted_dims = sorted(dims)
+    sorted_dims = sorted(active_dims)
     if num_out == 1:
         out_chunks[sorted_dims[0]] = out_sizes[0]
     else:
-        for dim, size in zip(dims, out_sizes):
+        for dim, size in zip(active_dims, out_sizes):
             # Note: when mapping 2D to 2D, this will be the only alteration to
             # out_chunks, the same as iris._lazy_data.map_complete_blocks
             out_chunks[dim] = size
@@ -370,10 +369,12 @@ def _map_complete_blocks(src, func, dims, out_sizes):
 
     return data.map_blocks(
         func,
+        *args,
         chunks=out_chunks,
         drop_axis=dropped_dims,
         new_axis=new_axis,
         dtype=src.dtype,
+        **kwargs,
     )
 
 
@@ -556,15 +557,6 @@ def _regrid_rectilinear_to_rectilinear__perform(src_cube, regrid_info, mdtol):
     grid_x, grid_y = regrid_info.target
     regridder = regrid_info.regridder
 
-    # Set up a function which can accept just chunk of data as an argument.
-    regrid = functools.partial(
-        _regrid_along_dims,
-        regridder,
-        dims=[grid_x_dim, grid_y_dim],
-        num_out_dims=2,
-        mdtol=mdtol,
-    )
-
     # Apply regrid to all the chunks of src_cube, ensuring first that all
     # chunks cover the entire horizontal plane (otherwise they would break
     # the regrid function).
@@ -575,9 +567,13 @@ def _regrid_rectilinear_to_rectilinear__perform(src_cube, regrid_info, mdtol):
         chunk_shape = grid_x.shape[::-1]
     new_data = _map_complete_blocks(
         src_cube,
-        regrid,
+        _regrid_along_dims,
         (grid_x_dim, grid_y_dim),
         chunk_shape,
+        regridder=regridder,
+        dims=[grid_x_dim, grid_y_dim],
+        num_out_dims=2,
+        mdtol=mdtol,
     )
 
     new_cube = _create_cube(
@@ -640,15 +636,6 @@ def _regrid_unstructured_to_rectilinear__perform(src_cube, regrid_info, mdtol):
     grid_x, grid_y = regrid_info.target
     regridder = regrid_info.regridder
 
-    # Set up a function which can accept just chunk of data as an argument.
-    regrid = functools.partial(
-        _regrid_along_dims,
-        regridder,
-        dims=[mesh_dim],
-        num_out_dims=2,
-        mdtol=mdtol,
-    )
-
     # Apply regrid to all the chunks of src_cube, ensuring first that all
     # chunks cover the entire horizontal plane (otherwise they would break
     # the regrid function).
@@ -659,9 +646,13 @@ def _regrid_unstructured_to_rectilinear__perform(src_cube, regrid_info, mdtol):
         chunk_shape = grid_x.shape[::-1]
     new_data = _map_complete_blocks(
         src_cube,
-        regrid,
+        _regrid_along_dims,
         (mesh_dim,),
         chunk_shape,
+        regridder=regridder,
+        dims=[mesh_dim],
+        num_out_dims=2,
+        mdtol=mdtol,
     )
 
     new_cube = _create_cube(
@@ -738,14 +729,6 @@ def _regrid_rectilinear_to_unstructured__perform(src_cube, regrid_info, mdtol):
     mesh, location = regrid_info.target
     regridder = regrid_info.regridder
 
-    # Set up a function which can accept just chunk of data as an argument.
-    regrid = functools.partial(
-        _regrid_along_dims,
-        regridder,
-        dims=[grid_x_dim, grid_y_dim],
-        num_out_dims=1,
-        mdtol=mdtol,
-    )
     if location == "face":
         face_node = mesh.face_node_connectivity
         # In face_node_connectivity: `location`= face, `connected` = node, so
@@ -761,9 +744,13 @@ def _regrid_rectilinear_to_unstructured__perform(src_cube, regrid_info, mdtol):
     # the regrid function).
     new_data = _map_complete_blocks(
         src_cube,
-        regrid,
+        _regrid_along_dims,
         (grid_x_dim, grid_y_dim),
         chunk_shape,
+        regridder=regridder,
+        dims=[grid_x_dim, grid_y_dim],
+        num_out_dims=1,
+        mdtol=mdtol,
     )
 
     new_cube = _create_cube(
@@ -836,13 +823,6 @@ def _regrid_unstructured_to_unstructured__perform(src_cube, regrid_info, mdtol):
     mesh, location = regrid_info.target
     regridder = regrid_info.regridder
 
-    regrid = functools.partial(
-        _regrid_along_dims,
-        regridder,
-        dims=[mesh_dim],
-        num_out_dims=1,
-        mdtol=mdtol,
-    )
     if location == "face":
         face_node = mesh.face_node_connectivity
         chunk_shape = (face_node.shape[face_node.location_axis],)
@@ -853,9 +833,13 @@ def _regrid_unstructured_to_unstructured__perform(src_cube, regrid_info, mdtol):
 
     new_data = _map_complete_blocks(
         src_cube,
-        regrid,
+        _regrid_along_dims,
         (mesh_dim,),
         chunk_shape,
+        regridder=regridder,
+        dims=[mesh_dim],
+        num_out_dims=1,
+        mdtol=mdtol,
     )
 
     new_cube = _create_cube(

noxfile.py

@@ -28,6 +28,9 @@
 #: GHA-CI environment variable hook.
 PY_VER = os.environ.get("PY_VER", ["3.10", "3.11", "3.12"])
 
+#: GHA-CI environment variable hook.
+COVERAGE = os.environ.get("COVERAGE", False)
+
 #: GHA-CI environment variable hook.
 #: If you change the IRIS_SOURCE here you will also need to change it in
 #: the tests, wheel and benchmark workflows.
@@ -302,8 +305,12 @@ def tests(session: nox.sessions.Session):
     # Install the esmf-regrid source in develop mode.
     session.install("--no-deps", "--editable", ".")
 
-    # Execute the tests.
-    session.run("pytest")
+    if COVERAGE:
+        # Execute the tests with code coverage.
+        session.run("pytest", "--cov-report=xml", "--cov")
+    else:
+        # Execute the tests.
+        session.run("pytest")
 
 
 @nox.session