Allow mesh partitioning

CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ## [Unreleased]
 
+### Added
+
+- [PR#686](https://github.com/SciTools/iris-esmf-regrid/pull/686)
+  Added the ability for ``Partition`` objects to work on meshes.
+  [@stephenworsley](https://github.com/stephenworsley)
+
 ## [0.15.0] - 2026-06-26
 
 ### Added

src/esmf_regrid/experimental/partition.py

@@ -1,6 +1,7 @@
 """Provides an interface for splitting up a large regridding task."""
 
 import esmpy
+from iris.experimental.mesh_coord_indexing import SETTING, Options
 import numpy as np
 
 from esmf_regrid.constants import Constants
@@ -17,7 +18,9 @@ def _get_chunk(cube, sl):
     full_slice = [np.s_[:]] * len(cube.shape)
     for s, d in zip(sl, grid_dims, strict=True):
         full_slice[d] = np.s_[s[0] : s[1]]
-    return cube[*full_slice]
+    with SETTING.context(Options.NEW_MESH):
+        result = cube[*full_slice]
+    return result
 
 
 def _determine_blocks(shape, chunks, num_chunks, explicit_blocks):
@@ -53,16 +56,11 @@ def _determine_blocks(shape, chunks, num_chunks, explicit_blocks):
                 raise ValueError(msg)
         bounds = [np.cumsum([0, *chunk]) for chunk in chunks]
         if len(bounds) == 1:
-            msg = "Chunks must have exactly two dimensions."
-            raise ValueError(msg)
-            # TODO: This is currently blocked by the fact that slicing an Iris cube on its mesh dimension
-            #  does not currently yield another cube with a mesh. When this is fixed, the following
-            #  code can be uncommented and the noqa on the following line can be removed.
-            # explicit_blocks = [
-            #     [[int(lower), int(upper)]]
-            #     for lower, upper in zip(bounds[0][:-1], bounds[0][1:], strict=True)
-            # ]
-        elif len(bounds) == 2:  # noqa: RET506
+            explicit_blocks = [
+                [[int(lower), int(upper)]]
+                for lower, upper in zip(bounds[0][:-1], bounds[0][1:], strict=True)
+            ]
+        elif len(bounds) == 2:
             explicit_blocks = [
                 [[int(ly), int(uy)], [int(lx), int(ux)]]
                 for ly, uy in zip(bounds[0][:-1], bounds[0][1:], strict=True)
@@ -159,12 +157,7 @@ def __init__(
         if src.mesh is None:
             grid_dims = _get_grid_dims(src)
         else:
-            msg = "Partition does not yet support source meshes."
-            raise NotImplementedError(msg)
-            # TODO: This is currently blocked by the fact that slicing an Iris cube on its mesh dimension
-            #  does not currently yield another cube with a mesh. When this is fixed, the following
-            #  code can be uncommented.
-            # grid_dims = (src.mesh_dim(),)
+            grid_dims = (src.mesh_dim(),)
         shape = tuple(src.shape[i] for i in grid_dims)
         self.tgt = tgt
         self.scheme = scheme

src/esmf_regrid/tests/unit/experimental/partition/test_Partition.py

@@ -113,19 +113,23 @@ def test_Partition_mesh_src(tmp_path):
     scheme = ESMFAreaWeighted(mdtol=1)
 
     src_chunks = (15000,)
-    with pytest.raises(NotImplementedError):
-        _ = Partition(src, tgt, scheme, files, src_chunks=src_chunks)
-
-    # TODO: when mesh partitioning becomes possible, uncomment.
-    # expected_src_chunks = [[[0, 15000]], [[15000, 30000]], [[30000, 45000]], [[45000, 60000]], [[60000, 75000]]]
-    # assert partition.src_blocks == expected_src_chunks
-    #
-    # partition.generate_files()
-    #
-    # result = partition.apply_regridders(src)
-    # expected = src.regrid(tgt, scheme)
-    # assert np.allclose(result.data, expected.data)
-    # assert result == expected
+    partition = Partition(src, tgt, scheme, files, src_chunks=src_chunks)
+
+    expected_src_chunks = [
+        [[0, 15000]],
+        [[15000, 30000]],
+        [[30000, 45000]],
+        [[45000, 60000]],
+        [[60000, 75000]],
+    ]
+    assert partition.src_blocks == expected_src_chunks
+
+    partition.generate_files()
+
+    result = partition.apply_regridders(src)
+    expected = src.regrid(tgt, scheme)
+    assert np.allclose(result.data, expected.data)
+    assert result == expected
 
 
 def test_Partition_curv_src(tmp_path):