cyclic_dataarray: extrapolate coords (regression) (#58)

* cyclic_dataarray: extrapolate coords (regression)

* changelog

Author: mathause

CHANGELOG.md

@@ -8,6 +8,10 @@
 
 ### Bug fixes
 
+- Fix a regression introduced in [#33](https://github.com/mathause/mplotutils/pull/33):
+  `cyclic_dataarray` now correctly extrapolates the coordinates
+  ([#58](https://github.com/mathause/mplotutils/pull/58)).
+
 ### Internal changes
 
 

mplotutils/cartopy_utils.py

@@ -112,7 +112,22 @@ def cyclic_dataarray(obj, coord="lon"):
 
     """
 
-    return obj.pad({coord: (0, 1)}, mode="wrap")
+    if coord not in obj.coords:
+        raise KeyError(f"Did not find '{coord}' in obj")
+
+    obj = obj.pad({coord: (0, 1)}, mode="wrap")
+
+    # extrapolate the coords
+    lon = obj[coord]
+
+    diff = lon.isel({coord: slice(None, -1)}).diff(coord)
+
+    if not np.allclose(diff, diff[0]):
+        raise ValueError(f"The coordinate '{coord}' must be equally spaced")
+
+    lon.data[-1] = lon[-2] + diff[0]
+
+    return obj.assign_coords({coord: lon})
 
 
 @_deprecate_positional_args("0.3")

mplotutils/tests/test_cyclic_dataarray.py

@@ -4,6 +4,29 @@
 from mplotutils import cyclic_dataarray
 
 
+@pytest.mark.parametrize("as_dataset", (True, False))
+def test_cyclic_dataarray_missing_coord(as_dataset):
+
+    da = xr.DataArray([1, 2, 3], dims=("x"), coords={"x": [0, 1, 2]}, name="data")
+    data = da.to_dataset() if as_dataset else da
+
+    with pytest.raises(KeyError, match="Did not find 'lon' in obj"):
+        cyclic_dataarray(data)
+
+    with pytest.raises(KeyError, match="Did not find 'longitude' in obj"):
+        cyclic_dataarray(data, "longitude")
+
+
+@pytest.mark.parametrize("as_dataset", (True, False))
+def test_cyclic_dataarray_not_equally_spaced(as_dataset):
+
+    da = xr.DataArray([1, 2, 3], dims=("x"), coords={"x": [0, 1, 2.1]}, name="data")
+    data = da.to_dataset() if as_dataset else da
+
+    with pytest.raises(ValueError, match=".*must be equally spaced"):
+        cyclic_dataarray(data, coord="x")
+
+
 @pytest.mark.parametrize("as_dataset", (True, False))
 def test_cyclic_dataarray(as_dataset):
 
@@ -14,7 +37,7 @@ def test_cyclic_dataarray(as_dataset):
 
     expected = [[1, 2, 3, 1], [4, 5, 6, 4]]
     da_expected = xr.DataArray(
-        expected, dims=("y", "x"), coords={"y": [1, 2], "x": [0, 1, 2, 0]}, name="data"
+        expected, dims=("y", "x"), coords={"y": [1, 2], "x": [0, 1, 2, 3]}, name="data"
     )
 
     data = da.to_dataset() if as_dataset else da