Dataset.apply works if func returns like-shaped ndarrays

xray/core/dataset.py

@@ -20,7 +20,7 @@
 from .alignment import align, partial_align
 from .coordinates import DatasetCoordinates, Indexes
 from .common import ImplementsDatasetReduce, AttrAccessMixin
-from .utils import Frozen, SortedKeysDict, ChainMap
+from .utils import Frozen, SortedKeysDict, ChainMap, maybe_wrap_array
 from .pycompat import iteritems, itervalues, basestring, OrderedDict
 
 
@@ -1452,8 +1452,9 @@ def apply(self, func, keep_attrs=False, args=(), **kwargs):
             Coordinates which are no longer used as the dimension of a
             noncoordinate are dropped.
         """
-        variables = OrderedDict((k, func(v, *args, **kwargs))
-                                for k, v in iteritems(self.data_vars))
+        variables = OrderedDict(
+            (k, maybe_wrap_array(v, func(v, *args, **kwargs)))
+            for k, v in iteritems(self.data_vars))
         attrs = self.attrs if keep_attrs else None
         return type(self)(variables, attrs=attrs)
 

xray/core/groupby.py

@@ -6,7 +6,7 @@
 from .alignment import concat
 from .common import ImplementsArrayReduce, ImplementsDatasetReduce
 from .pycompat import zip
-from .utils import peek_at
+from .utils import peek_at, maybe_wrap_array
 from .variable import Variable, Coordinate
 
 
@@ -242,12 +242,6 @@ def apply(self, func, shortcut=False, **kwargs):
         applied : DataArray
             The result of splitting, applying and combining this array.
         """
-        def maybe_wrap_array(arr, f):
-            # in case func lost array's metadata
-            if isinstance(f, np.ndarray) and f.shape == arr.shape:
-                return arr.__array_wrap__(f)
-            else:
-                return f
         if shortcut:
             grouped = self._iter_grouped_shortcut()
         else:

xray/core/utils.py

@@ -96,6 +96,19 @@ def safe_cast_to_index(array):
     return index
 
 
+def maybe_wrap_array(original, new_array):
+    """Wrap a transformed array with __array_wrap__ is it can be done safely.
+
+    This lets us treat arbitrary functions that take and return ndarray objects
+    like ufuncs, as long as they return an array with the same shape.
+    """
+    # in case func lost array's metadata
+    if isinstance(new_array, np.ndarray) and new_array.shape == original.shape:
+        return original.__array_wrap__(new_array)
+    else:
+        return new_array
+
+
 def equivalent(first, second):
     """Compare two objects for equivalence (identity or equality), using
     array_equiv if either object is an ndarray

xray/test/test_dataset.py

@@ -1514,6 +1514,10 @@ def scale(x, multiple=1):
         self.assertDataArrayEqual(actual['var1'], 2 * data['var1'])
         self.assertDataArrayIdentical(actual['numbers'], data['numbers'])
 
+        actual = data.apply(np.asarray)
+        expected = data.drop_vars('time') # time is not used on a data var
+        self.assertDatasetEqual(expected, actual)
+
     def make_example_math_dataset(self):
         variables = OrderedDict(
             [('bar', ('x', np.arange(100, 400, 100))),