Spherical padding and faster tests

src/xarray_regrid/methods/conservative.py

@@ -72,8 +72,10 @@ def conservative_regrid(
 
     # Make sure the regridding coordinates are sorted
     coord_names = [coord for coord in target_ds.coords if coord in data.coords]
-    target_ds_sorted = target_ds.sortby(coord_names)
-    data = data.sortby(list(coord_names))
+    target_ds_sorted = target_ds.copy()
+    for coord_name in coord_names:
+        target_ds_sorted = utils.ensure_monotonic(target_ds_sorted, coord_name)
+        data = utils.ensure_monotonic(data, coord_name)
     coords = {name: target_ds_sorted[name] for name in coord_names}
 
     regridded_data = utils.call_on_dataset(

src/xarray_regrid/utils.py

@@ -198,24 +198,6 @@ def common_coords(
     return sorted([str(coord) for coord in coords])
 
 
-@overload
-def call_on_dataset(
-    func: Callable[..., xr.Dataset],
-    obj: xr.DataArray,
-    *args: Any,
-    **kwargs: Any,
-) -> xr.DataArray: ...
-
-
-@overload
-def call_on_dataset(
-    func: Callable[..., xr.Dataset],
-    obj: xr.Dataset,
-    *args: Any,
-    **kwargs: Any,
-) -> xr.Dataset: ...
-
-
 def call_on_dataset(
     func: Callable[..., xr.Dataset],
     obj: xr.DataArray | xr.Dataset,
@@ -266,8 +248,8 @@ def format_for_regrid(
     # Apply formatting
     for coord_type, coord in formatted_coords.items():
         # Make sure formatted coords are sorted
-        obj = obj.sortby(coord)
-        target = target.sortby(coord)
+        obj = ensure_monotonic(obj, coord)
+        target = ensure_monotonic(target, coord)
         obj = coord_handlers[coord_type]["func"](obj, target, formatted_coords)
         # Coerce back to a single chunk if that's what was passed
         if len(orig_chunksizes.get(coord, [])) == 1:
@@ -306,21 +288,24 @@ def format_lat(
     # Only pad if global but don't have edge values directly at poles
     # NOTE: could use xr.pad here instead of xr.concat, but none of the
     # modes are an exact fit for this scheme
+    lat_vals = obj.coords[lat_coord].values
     # South pole
     if dy - polar_lat >= obj.coords[lat_coord].values[0] > -polar_lat:
         south_pole = obj.isel({lat_coord: 0})
         if lon_coord is not None:
             south_pole = south_pole.mean(lon_coord)
         obj = xr.concat([south_pole, obj], dim=lat_coord)  # type: ignore
-        obj.coords[lat_coord].values[0] = -polar_lat
+        lat_vals = np.concatenate([[-polar_lat], lat_vals])
 
     # North pole
     if polar_lat - dy <= obj.coords[lat_coord].values[-1] < polar_lat:
         north_pole = obj.isel({lat_coord: -1})
         if lon_coord is not None:
             north_pole = north_pole.mean(lon_coord)
         obj = xr.concat([obj, north_pole], dim=lat_coord)  # type: ignore
-        obj.coords[lat_coord].values[-1] = polar_lat
+        lat_vals = np.concatenate([lat_vals, [polar_lat]])
+
+    obj = update_coord(obj, lat_coord, lat_vals)
 
     return obj
 
@@ -349,11 +334,9 @@ def format_lon(
         source_vals < wrap_point - 360, source_vals + 360, source_vals
     )
     source_vals = np.where(source_vals > wrap_point, source_vals - 360, source_vals)
-    obj.coords[lon_coord].values[:] = source_vals
+    obj = update_coord(obj, lon_coord, source_vals)
 
-    # Shift operations can produce duplicates
-    # Simplest solution is to drop them and add back when padding
-    obj = obj.sortby(lon_coord).drop_duplicates(lon_coord)
+    obj = ensure_monotonic(obj, lon_coord)
 
     # Only pad if domain is global in lon
     source_lon = obj.coords[lon_coord]
@@ -368,14 +351,12 @@ def format_lon(
         left_pad = int(np.ceil(np.max([left_pad, 0])))
         right_pad = int(np.ceil(np.max([right_pad, 0])))
         obj = obj.pad({lon_coord: (left_pad, right_pad)}, mode="wrap", keep_attrs=True)
+        lon_vals = obj.coords[lon_coord].values
         if left_pad:
-            obj.coords[lon_coord].values[:left_pad] = (
-                source_lon.values[-left_pad:] - 360
-            )
+            lon_vals[:left_pad] = source_lon.values[-left_pad:] - 360
         if right_pad:
-            obj.coords[lon_coord].values[-right_pad:] = (
-                source_lon.values[:right_pad] + 360
-            )
+            lon_vals[-right_pad:] = source_lon.values[:right_pad] + 360
+        obj = update_coord(obj, lon_coord, lon_vals)
 
     return obj
 
@@ -388,3 +369,50 @@ def coord_is_covered(
     left_covered = obj[coord].min() <= target[coord].min() - pad
     right_covered = obj[coord].max() >= target[coord].max() + pad
     return bool(left_covered.item() and right_covered.item())
+
+
+@overload
+def ensure_monotonic(obj: xr.DataArray, coord: Hashable) -> xr.DataArray: ...
+
+
+@overload
+def ensure_monotonic(obj: xr.Dataset, coord: Hashable) -> xr.Dataset: ...
+
+
+def ensure_monotonic(
+    obj: xr.DataArray | xr.Dataset, coord: Hashable
+) -> xr.DataArray | xr.Dataset:
+    """Ensure that an object has monotonically increasing indexes for a
+    given coordinate. Only sort and drop duplicates if needed because this
+    requires reindexing which can be expensive."""
+    is_sorted = (obj.coords[coord].diff(coord) >= 0).all().compute().item()
+    if not is_sorted:
+        obj = obj.sortby(coord)
+    has_duplicates = (
+        np.unique(obj.coords[coord].values).size < obj.coords[coord].values.size
+    )
+    if has_duplicates:
+        obj = obj.drop_duplicates(coord)
+    return obj
+
+
+@overload
+def update_coord(
+    obj: xr.DataArray, coord: Hashable, coord_vals: np.ndarray
+) -> xr.DataArray: ...
+
+
+@overload
+def update_coord(
+    obj: xr.Dataset, coord: Hashable, coord_vals: np.ndarray
+) -> xr.Dataset: ...
+
+
+def update_coord(
+    obj: xr.DataArray | xr.Dataset, coord: Hashable, coord_vals: np.ndarray
+) -> xr.DataArray | xr.Dataset:
+    """Update the values of a coordinate, ensuring indexes stay in sync."""
+    attrs = obj.coords[coord].attrs
+    obj = obj.assign_coords({coord: coord_vals})
+    obj.coords[coord].attrs = attrs
+    return obj