make round tripping more thorough

Author: stephenworsley

pyproject.toml

@@ -298,5 +298,5 @@ convention = "numpy"
 [tool.ruff.lint.pylint]
 # TODO: refactor to reduce complexity, if possible
 max-args = 10
-max-branches = 23
+max-branches = 25
 max-statements = 110

src/esmf_regrid/experimental/io.py

@@ -20,7 +20,6 @@
     ESMFAreaWeightedRegridder,
     ESMFBilinear,
     ESMFBilinearRegridder,
-    ESMFNearest,
     ESMFNearestRegridder,
     GridRecord,
     MeshRecord,
@@ -76,6 +75,7 @@
     "extrap_method": _EXTRAP_METHOD_DICT,
     "unmapped_action": _UNMAPPED_ACTION_DICT,
 }
+_ESMF_BOOL_ARGS = ["ignore_degenerate", "large_file"]
 
 
 def _add_mask_to_cube(mask, cube, name):
@@ -294,7 +294,7 @@ def save_regridder(rg, filename, allow_partial=False):
         if tgt_slice is None:
             tgt_slice = []
         tgt_slice_cube = Cube(
-            src_slice, long_name=_TGT_SLICE_NAME, var_name=_TGT_SLICE_NAME
+            tgt_slice, long_name=_TGT_SLICE_NAME, var_name=_TGT_SLICE_NAME
         )
         extra_cubes = [src_slice_cube, tgt_slice_cube]
 
@@ -416,11 +416,11 @@ def load_regridder(filename, allow_partial=False):
     mdtol = weights_cube.attributes[_MDTOL]
 
     if src_cube.coords(_SOURCE_MASK_NAME):
-        use_src_mask = src_cube.coord(_SOURCE_MASK_NAME).points
+        use_src_mask = src_cube.coord(_SOURCE_MASK_NAME).points.astype(bool)
     else:
         use_src_mask = False
     if tgt_cube.coords(_TARGET_MASK_NAME):
-        use_tgt_mask = tgt_cube.coord(_TARGET_MASK_NAME).points
+        use_tgt_mask = tgt_cube.coord(_TARGET_MASK_NAME).points.astype(bool)
     else:
         use_tgt_mask = False
 
@@ -433,6 +433,9 @@ def load_regridder(filename, allow_partial=False):
     for arg, arg_dict in _ESMF_ENUM_ARGS.items():
         if arg in esmf_args:
             esmf_args[arg] = arg_dict[esmf_args[arg]]
+    for arg in _ESMF_BOOL_ARGS:
+        if arg in esmf_args:
+            esmf_args[arg] = bool(esmf_args[arg])
 
     if scheme is GridToMeshESMFRegridder:
         resolution_keyword = _SOURCE_RESOLUTION
@@ -463,16 +466,21 @@ def load_regridder(filename, allow_partial=False):
             Constants.Method.BILINEAR: ESMFBilinear,
         }[method]
         mdtol = kwargs.pop(_MDTOL, None)
-        mdtol_dict = {}
+        sub_kwargs = {}
         if mdtol is not None:
-            mdtol_dict[_MDTOL] = mdtol
+            sub_kwargs[_MDTOL] = mdtol
         regridder = scheme(
             src_cube,
             tgt_cube,
             src_slice,
             tgt_slice,
             weight_matrix,
-            sub_scheme(**mdtol_dict),
+            sub_scheme(
+                use_src_mask=use_src_mask,
+                use_tgt_mask=use_tgt_mask,
+                esmf_args=esmf_args,
+                **sub_kwargs,
+            ),
             **kwargs,
         )
     else:

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

@@ -1,7 +1,10 @@
 """Unit tests for :mod:`esmf_regrid.experimental.partition`."""
 
+import numpy as np
+
 from esmf_regrid import ESMFAreaWeighted
 from esmf_regrid.experimental._partial import PartialRegridder
+from esmf_regrid.experimental.io import load_regridder, save_regridder
 from esmf_regrid.tests.unit.schemes.test__cube_to_GridInfo import (
     _grid_cube,
 )
@@ -23,3 +26,26 @@ def test_PartialRegridder_repr():
         "scheme=ESMFAreaWeighted(mdtol=0.5, use_src_mask=False, use_tgt_mask=False, esmf_args={}))"
     )
     assert repr(pr) == expected_repr
+
+
+def test_PartialRegridder_roundtrip(tmp_path):
+    """Test load/save for PartialRegridder instance."""
+    src = _grid_cube(10, 15, (-180, 180), (-90, 90), circular=True)
+    mask = np.zeros_like(src.data)
+    mask[0, 0] = 1
+    src.data = np.ma.array(src.data, mask=mask)
+    tgt = _grid_cube(5, 10, (-180, 180), (-90, 90), circular=True)
+    src_slice = [[10, 20], [15, 30]]
+    tgt_slice = [[0, 5], [0, 10]]
+    weights = None
+    scheme = ESMFAreaWeighted(
+        mdtol=0.5, use_src_mask=src.data.mask, esmf_args={"ignore_degenerate": True}
+    )
+
+    pr = PartialRegridder(src, tgt, src_slice, tgt_slice, weights, scheme)
+    file = tmp_path / "partial.nc"
+
+    save_regridder(pr, file, allow_partial=True)
+    loaded_pr = load_regridder(file, allow_partial=True)
+
+    assert repr(loaded_pr) == repr(pr)