apply to dataset

doc/api.rst

@@ -19,6 +19,7 @@ Top-level functions
    apply_ufunc
    align
    broadcast
+   call_on_dataset
    concat
    merge
    combine_by_coords

doc/whats-new.rst

@@ -64,6 +64,9 @@ New Features
   :py:class:`~core.groupby.DataArrayGroupBy`, inspired by pandas'
   :py:meth:`~pandas.core.groupby.GroupBy.get_group`.
   By `Deepak Cherian <https://github.com/dcherian>`_.
+- Add :py:func:`call_on_dataset` as a way to apply functions expecting
+  :py:class:`Dataset` objects to :py:class:`DataArray` objects (:issue:`4837`, :pull:`4863`).
+  By `Justus Magin <https://github.com/keewis>`_.
 - Add a ``combine_attrs`` parameter to :py:func:`open_mfdataset` (:pull:`4971`).
   By `Justus Magin <https://github.com/keewis>`_.
 - Disable the `cfgrib` backend if the `eccodes` library is not installed (:pull:`5083`). By `Baudouin Raoult <https://github.com/b8raoult>`_.

xarray/__init__.py

@@ -18,7 +18,15 @@
 from .core.alignment import align, broadcast
 from .core.combine import combine_by_coords, combine_nested
 from .core.common import ALL_DIMS, full_like, ones_like, zeros_like
-from .core.computation import apply_ufunc, corr, cov, dot, polyval, where
+from .core.computation import (
+    apply_ufunc,
+    call_on_dataset,
+    corr,
+    cov,
+    dot,
+    polyval,
+    where,
+)
 from .core.concat import concat
 from .core.dataarray import DataArray
 from .core.dataset import Dataset
@@ -46,6 +54,7 @@
     # Top-level functions
     "align",
     "apply_ufunc",
+    "call_on_dataset",
     "as_variable",
     "broadcast",
     "cftime_range",

xarray/core/computation.py

@@ -1151,6 +1151,99 @@ def earth_mover_distance(first_samples,
         return apply_array_ufunc(func, *args, dask=dask)
 
 
+def call_on_dataset(func, obj, *args, **kwargs):
+    """apply a function expecting a Dataset to a xarray object
+
+    Parameters
+    ----------
+    func : callable
+        A function expecting a Dataset as its first parameter.
+    obj : DataArray or Dataset
+        The dataset to apply ``func`` to. If a ``DataArray``, convert it to a single
+        variable ``Dataset`` first.
+    *args, **kwargs
+        Additional arguments to ``func``
+
+    Returns
+    -------
+    DataArray or Dataset
+        The result of ``func(obj, *args, **kwargs)`` with the same type as ``obj``.
+
+    Notes
+    -----
+    DataArray objects without a name (or named ``None``) will be renamed to
+    ``"<this-array>"`` before being passed to ``func``. The empty name will be restored
+    for the result of the call.
+
+    See Also
+    --------
+    Dataset.map
+    Dataset.pipe
+    DataArray.pipe
+
+    Examples
+    --------
+    >>> def f(ds):
+    ...     return xr.Dataset(
+    ...         {
+    ...             name: var * var.attrs.get("scale", 1)
+    ...             for name, var in ds.data_vars.items()
+    ...         },
+    ...         coords=ds.coords,
+    ...         attrs=ds.attrs,
+    ...     )
+    ...
+    >>> ds = xr.Dataset(
+    ...     {"a": ("x", [3, 4], {"scale": 0.5}), "b": ("x", [-1, 1], {"scale": 1.5})},
+    ...     coords={"x": [0, 1]},
+    ...     attrs={"attr": "value"},
+    ... )
+    >>> ds
+    <xarray.Dataset>
+    Dimensions:  (x: 2)
+    Coordinates:
+      * x        (x) int64 0 1
+    Data variables:
+        a        (x) int64 3 4
+        b        (x) int64 -1 1
+    Attributes:
+        attr:     value
+    >>> xr.call_on_dataset(f, ds)
+    <xarray.Dataset>
+    Dimensions:  (x: 2)
+    Coordinates:
+      * x        (x) int64 0 1
+    Data variables:
+        a        (x) float64 1.5 2.0
+        b        (x) float64 -1.5 1.5
+    Attributes:
+        attr:     value
+    >>> xr.call_on_dataset(f, ds.a)
+    <xarray.DataArray 'a' (x: 2)>
+    array([1.5, 2. ])
+    Coordinates:
+      * x        (x) int64 0 1
+    """
+    from .dataarray import _THIS_ARRAY, DataArray
+    from .parallel import dataarray_to_dataset, dataset_to_dataarray
+
+    if isinstance(obj, DataArray):
+        ds = dataarray_to_dataset(obj)
+        if obj.name is None:
+            ds = ds.rename({_THIS_ARRAY: "<this-array>"})
+    else:
+        ds = obj
+
+    result = func(ds, *args, **kwargs)
+
+    if isinstance(obj, DataArray):
+        if obj.name is None:
+            result = result.rename({"<this-array>": None})
+        result = dataset_to_dataarray(result)
+
+    return result
+
+
 def cov(da_a, da_b, dim=None, ddof=1):
     """
     Compute covariance between two DataArray objects along a shared dimension.

xarray/tests/test_computation.py

@@ -468,6 +468,50 @@ def test_apply_groupby_add():
         add(data_array.groupby("y"), data_array.groupby("x"))
 
 
+@pytest.mark.parametrize(
+    ["obj", "expected"],
+    (
+        pytest.param(
+            xr.DataArray(
+                [0, 1],
+                coords={
+                    "x": ("x", [-1, 1], {"a": 1, "b": 2}),
+                    "u": ("x", [2, 3], {"c": 3}),
+                },
+                dims="x",
+                attrs={"d": 4, "e": 5},
+            ),
+            xr.DataArray([0, 1], coords={"x": [-1, 1], "u": ("x", [2, 3])}, dims="x"),
+            id="DataArray",
+        ),
+        pytest.param(
+            xr.Dataset(
+                {"a": ("x", [1, 2], {"a": 1, "b": 2}), "b": ("x", [0, 1], {"c": 3})},
+                coords={
+                    "x": ("x", [-1, 1], {"d": 4, "e": 5}),
+                    "u": ("x", [2, 3], {"f": 6}),
+                },
+            ),
+            xr.Dataset(
+                {"a": ("x", [1, 2]), "b": ("x", [0, 1])},
+                coords={"x": [-1, 1], "u": ("x", [2, 3])},
+            ),
+            id="Dataset",
+        ),
+    ),
+)
+def test_call_on_dataset(obj, expected):
+    def clear_all_attrs(ds):
+        new_ds = ds.copy()
+        for var in new_ds.variables.values():
+            var.attrs.clear()
+        new_ds.attrs.clear()
+        return new_ds
+
+    actual = xr.call_on_dataset(clear_all_attrs, obj)
+    assert_identical(actual, expected)
+
+
 def test_unified_dim_sizes():
     assert unified_dim_sizes([xr.Variable((), 0)]) == {}
     assert unified_dim_sizes([xr.Variable("x", [1]), xr.Variable("x", [1])]) == {"x": 1}