refactor cyclic_dataarray (#33)

* refactor cyclic_dataarray

* changelog

Author: mathause

CHANGELOG.md

@@ -19,6 +19,7 @@
    gridspec.
  * Replaced `ax.get_subplotspec().get_geometry()` with `ax.get_subplotspec().get_geometry()`
    as the former was deprecated in matplotlib (#8).
+ * Refactor `mpu.cyclic_dataarray` using `obj.pad` ([#33](https://github.com/mathause/mplotutils/pull/33)).
  * Enabled CI on github actions (#9).
  * Formatted with black and isort, checked with flake8.
 

licenses/PLOT_ALL_IN_NCFILE_LICENSE

@@ -1,9 +1,6 @@
-# code by M.Hauser
-
 import warnings
 
 import cartopy.crs as ccrs
-import cartopy.util as cutil
 import matplotlib.pyplot as plt
 import numpy as np
 import shapely.geometry as sgeom
@@ -127,14 +124,29 @@ def _is_monotonic(coord, axis=0):
 # =============================================================================
 
 
-def cyclic_dataarray(da, coord="lon"):
-    """
-    Add a cyclic coordinate point to a DataArray along a specified named dimension.
+def cyclic_dataarray(obj, coord="lon"):
+    """Add a cyclic coordinate point to a DataArray or Dataset along a dimension.
+
+    Parameters
+    ----------
+    obj : xr.Dataset | xr.DataArray
+        Object to add the cyclic data point to.
+    coord : str, default: "lon"
+        Name of the
 
+    Returns
+    -------
+    obj_cyclic : xr.Dataset | xr.DataArray
+        The same as `obj` with a cyclic data point added.
+
+    Examples
+    --------
     >>> import xarray as xr
-    >>> data = xr.DataArray([[1, 2, 3], [4, 5, 6]],
-    ...                      coords={'x': [1, 2], 'y': range(3)},
-    ...                      dims=['x', 'y'])
+    >>> data = xr.DataArray(
+    ...     [[1, 2, 3], [4, 5, 6]],
+    ...     coords={'x': [1, 2], 'y': range(3)},
+    ...     dims=['x', 'y']
+    ... )
     >>> data_cyclic = cyclic_dataarray(data, 'y')
     >>> data_cyclic
     <xarray.DataArray (x: 2, y: 4)>
@@ -144,38 +156,9 @@ def cyclic_dataarray(da, coord="lon"):
       * x        (x) int64 1 2
       * y        (y) int64 0 1 2 3
 
-    Notes
-    -----
-    After: https://github.com/darothen/plot-all-in-ncfile/blob/master/plot_util.py
-
     """
-    import xarray as xr
 
-    assert isinstance(da, xr.DataArray)
-
-    lon_idx = da.dims.index(coord)
-    cyclic_data, cyclic_coord = cutil.add_cyclic_point(
-        da.values, coord=da.coords[coord], axis=lon_idx
-    )
-
-    # Copy and add the cyclic coordinate and data
-    new_coords = dict(da.coords)
-    new_coords[coord] = cyclic_coord
-    new_values = cyclic_data
-
-    new_da = xr.DataArray(new_values, dims=da.dims, coords=new_coords)
-
-    # Copy the attributes for the re-constructed data and coords
-    for att, val in da.attrs.items():
-        new_da.attrs[att] = val
-    for c in da.coords:
-        for att in da.coords[c].attrs:
-            new_da.coords[c].attrs[att] = da.coords[c].attrs[att]
-
-    return new_da
-
-
-# =============================================================================
+    return obj.pad({coord: (0, 1)}, mode="wrap")
 
 
 def ylabel_map(s, labelpad=None, size=None, weight=None, y=0.5, ax=None, **kwargs):

mplotutils/cartopy_utils.py

@@ -1,28 +1,31 @@
-import numpy as np
+import pytest
 import xarray as xr
 
-from mplotutils.cartopy_utils import cyclic_dataarray
+from mplotutils import cyclic_dataarray
 
 
-def test_cyclic_dataarray():
+@pytest.mark.parametrize("as_dataset", (True, False))
+def test_cyclic_dataarray(as_dataset):
 
-    data = xr.DataArray(
-        [[1, 2, 3], [4, 5, 6]], coords={"x": [1, 2], "y": range(3)}, dims=["x", "y"]
+    data = [[1, 2, 3], [4, 5, 6]]
+    da = xr.DataArray(
+        data, dims=("y", "x"), coords={"y": [1, 2], "x": [0, 1, 2]}, name="data"
     )
 
-    res = cyclic_dataarray(data, "y")
-
-    expected_data = np.asarray([[1, 2, 3, 1], [4, 5, 6, 4]])
+    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"
+    )
 
-    np.testing.assert_allclose(res, expected_data)
+    data = da.to_dataset() if as_dataset else da
+    expected = da_expected.to_dataset() if as_dataset else da_expected
 
-    np.testing.assert_allclose(res.y, [0, 1, 2, 3])
-    np.testing.assert_allclose(res.x, [1, 2])
+    result = cyclic_dataarray(data, "x")
+    xr.testing.assert_identical(result, expected)
 
     # per default use 'lon'
-    data = xr.DataArray(
-        [[1, 2, 3], [4, 5, 6]], coords={"x": [1, 2], "lon": range(3)}, dims=["x", "lon"]
-    )
+    data = data.rename(x="lon")
+    expected = expected.rename(x="lon")
 
-    res = cyclic_dataarray(data)
-    np.testing.assert_allclose(res, expected_data)
+    result = cyclic_dataarray(data)
+    xr.testing.assert_identical(result, expected)