add bilinear regridding support

Author: stephenworsley

src/esmf_regrid/experimental/partition.py

@@ -1,11 +1,12 @@
 """Provides an interface for splitting up a large regridding task."""
-
+import esmpy
 import numpy as np
 
 from esmf_regrid.constants import Constants
 from esmf_regrid.experimental._partial import PartialRegridder
 from esmf_regrid.experimental.io import load_regridder, save_regridder
 from esmf_regrid.schemes import _get_grid_dims
+from pyarrow import dictionary
 
 
 def _get_chunk(cube, sl):
@@ -132,6 +133,12 @@ def __init__(
         if scheme._method == Constants.Method.NEAREST:
             msg = "The `Nearest` method is not implemented."
             raise NotImplementedError(msg)
+        if scheme._method == Constants.Method.BILINEAR:
+            pole_method = scheme.esmf_args.get("pole_method")
+            if pole_method != esmpy.PoleMethod.NONE:
+                msg = ("Bilinear regridding must have a `pole_method` of `esmpy.PoleMethod.NONE` in "
+                       "the `esmf_args` in order for Partition to work.`")
+                raise ValueError(msg)
         # TODO: Extract a slice of the cube.
         self.src = src
         if src.mesh is None:

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

@@ -1,10 +1,11 @@
 """Unit tests for :mod:`esmf_regrid.experimental.partition`."""
 
 import dask.array as da
+import esmpy
 import numpy as np
 import pytest
 
-from esmf_regrid import ESMFAreaWeighted, ESMFNearest
+from esmf_regrid import ESMFAreaWeighted, ESMFBilinear, ESMFNearest
 from esmf_regrid.experimental.partition import Partition
 from esmf_regrid.tests.unit.schemes.test__cube_to_GridInfo import (
     _curvilinear_cube,
@@ -155,6 +156,32 @@ def test_Partition_curv_src(tmp_path):
     assert result == expected
 
 
+def test_Partition_bilinear(tmp_path):
+    """Test Partition class for bilinear regridding."""
+    src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True)
+    src.data = np.arange(150 * 500).reshape([500, 150])
+    tgt = _grid_cube(16, 36, (-180, 180), (-90, 90), circular=True)
+
+    files = [tmp_path / f"partial_{x}.nc" for x in range(5)]
+    src_chunks = (100, 150)
+
+    bad_scheme = ESMFBilinear()
+    with pytest.raises(ValueError):
+        _ = Partition(src, tgt, bad_scheme, files, src_chunks=src_chunks)
+
+    # The pole_method must be NONE for bilinear regridding partitions to work.
+    scheme = ESMFBilinear(esmf_args={"pole_method": esmpy.PoleMethod.NONE})
+
+    partition = Partition(src, tgt, scheme, files, src_chunks=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_mesh_tgt(tmp_path):
     """Test Partition class when the target has a mesh."""
     src = _grid_cube(150, 500, (-180, 180), (-90, 90), circular=True)