handle variable chunks individually and test

Author: slevang

src/xarray_regrid/utils.py

@@ -249,8 +249,6 @@ def format_for_regrid(
     Currently handles padding of spherical geometry if lat/lon coordinates can
     be inferred and the domain size requires boundary padding.
     """
-    orig_chunksizes = obj.chunksizes
-
     # Special-cased coordinates with accepted names and formatting function
     coord_handlers: dict[str, CoordHandler] = {
         "lat": {"names": ["lat", "latitude"], "func": format_lat},
@@ -270,15 +268,22 @@ def format_for_regrid(
                 formatted_coords[coord_type] = str(coord)
 
     # Apply formatting
+    result = obj.copy()
     for coord_type, coord in formatted_coords.items():
         # Make sure formatted coords are sorted
-        obj = ensure_monotonic(obj, coord)
+        result = ensure_monotonic(result, coord)
         target = ensure_monotonic(target, coord)
-        obj = coord_handlers[coord_type]["func"](obj, target, formatted_coords)
+        result = coord_handlers[coord_type]["func"](result, target, formatted_coords)
+
         # Coerce back to a single chunk if that's what was passed
-        if len(orig_chunksizes.get(coord, [])) == 1:
-            obj = obj.chunk({coord: -1})
-    return obj
+        if isinstance(obj, xr.DataArray) and len(obj.chunksizes.get(coord, ())) == 1:
+            result = result.chunk({coord: -1})
+        elif isinstance(obj, xr.Dataset):
+            for var in result.data_vars:
+                if len(obj[var].chunksizes.get(coord, ())) == 1:
+                    result[var] = result[var].chunk({coord: -1})
+
+    return result
 
 
 def format_lat(

tests/test_format.py

@@ -216,3 +216,40 @@ def test_stats():
     # And preserve integer dtypes
     assert formatted.data.dtype == source.data.dtype
     assert (formatted.longitude.diff("longitude") == 1).all()
+
+
+def test_maintain_single_chunk():
+    dx_source = 2
+    source = xarray_regrid.Grid(
+        north=90 - dx_source / 2,
+        east=360 - dx_source / 2,
+        south=-90 + dx_source / 2,
+        west=0 + dx_source / 2,
+        resolution_lat=dx_source,
+        resolution_lon=dx_source,
+    ).create_regridding_dataset()
+    source["a"] = xr.DataArray(
+        np.ones((source.latitude.size, source.longitude.size)),
+        dims=["latitude", "longitude"],
+        coords={"latitude": source.latitude, "longitude": source.longitude},
+    ).chunk({"latitude": -1, "longitude": -1})
+    source["b"] = source.a.copy().chunk({"latitude": 45, "longitude": 90})
+
+    dx_target = 1
+    target = xarray_regrid.Grid(
+        north=90,
+        east=360,
+        south=-90,
+        west=0,
+        resolution_lat=dx_target,
+        resolution_lon=dx_target,
+    ).create_regridding_dataset()
+
+    # dataset
+    formatted = format_for_regrid(source, target)
+    assert formatted.a.chunks == ((92,), (182,))
+    assert formatted.b.chunks == ((1, 45, 45, 1), (1, 90, 90, 1))
+
+    # dataarray
+    formatted = format_for_regrid(source.a, target)
+    assert formatted.chunks == ((92,), (182,))