Skip to content

Commit 98404ad

Browse files
committed
ENH: Update units in .cf.differentiate
1 parent 24dd81b commit 98404ad

File tree

2 files changed

+31
-1
lines changed

2 files changed

+31
-1
lines changed

cf_xarray/accessor.py

+20
Original file line numberDiff line numberDiff line change
@@ -2102,6 +2102,26 @@ def differentiate(
21022102
(_single(_get_coords),), self._obj, coord, error=False, default=[coord]
21032103
)[0]
21042104
result = self._obj.differentiate(coord, *xr_args, **xr_kwargs)
2105+
if isinstance(self._obj, DataArray):
2106+
try:
2107+
result.attrs["units"] = "{:s} / ({:s})".format(
2108+
self._obj.attrs["units"], self._obj[coord].attrs["units"]
2109+
)
2110+
except KeyError:
2111+
pass
2112+
else:
2113+
try:
2114+
coord_units = self._obj[coord].attrs["units"]
2115+
except KeyError:
2116+
pass
2117+
else:
2118+
for name in result.data_vars:
2119+
try:
2120+
result[name].attrs["units"] = "{:s} / ({:s})".format(
2121+
self._obj[name].attrs["units"], coord_units
2122+
)
2123+
except KeyError:
2124+
pass
21052125
if positive_upward:
21062126
coord = self._obj[coord]
21072127
attrs = coord.attrs

cf_xarray/tests/test_accessor.py

+11-1
Original file line numberDiff line numberDiff line change
@@ -1472,12 +1472,22 @@ def test_rename_tuple():
14721472

14731473
@pytest.mark.parametrize("ds", datasets)
14741474
def test_differentiate(ds):
1475+
from ..units import units
14751476
# Add data_var and coord to test _get_coords
14761477
ds["lon_var"] = ds["lon"]
14771478
ds = ds.assign_coords(lon_coord=ds["lon"])
14781479

1480+
result = ds.cf.differentiate("lon")
1481+
lon_units = units.Unit(ds["lon"].attrs["units"])
1482+
for name in ds.data_vars:
1483+
expected_units = (
1484+
units.Unit(ds[name].attrs["units"]) / lon_units
1485+
)
1486+
assert units.Unit(result[name].attrs["units"]) == expected_units
1487+
del result[name].attrs["units"]
1488+
14791489
# Coordinate
1480-
assert_identical(ds.differentiate("lon"), ds.cf.differentiate("lon"))
1490+
assert_identical(ds.differentiate("lon"), result)
14811491

14821492
# Multiple coords (test error raised by _single)
14831493
with pytest.raises(KeyError, match=".*I expected only one."):

0 commit comments

Comments
 (0)