Repository: pydata/xarray

Issue:
Inconsistent signature of `apply`
`Dataset.apply` allows passing positional arguments to `func` as `args=()`, whereas `DataArrayGroupBy.apply` and `DatasetGroupBy.apply` do not. The documentation also incorrectly suggests that the reverse is true:

`Dataset.apply`:
https://github.com/pydata/xarray/blob/4de8dbc3b1de461c0c9d3b002e55d60b46d2e6d2/xarray/core/dataset.py#L2832-L2840

`DataArrayGroupBy.apply`:
https://github.com/pydata/xarray/blob/4de8dbc3b1de461c0c9d3b002e55d60b46d2e6d2/xarray/core/groupby.py#L473-L478

`DatasetGroupBy.apply`:
https://github.com/pydata/xarray/blob/4de8dbc3b1de461c0c9d3b002e55d60b46d2e6d2/xarray/core/groupby.py#L580-L585

I feel that `func`s passed to `GroupBy` objects are generally equally deserving of positional arguments, and I don't see a reason why this isn't the case.

Discussion:
shoyer: Looks good to me, thanks!

For consistency, can you update the apply method on the classes in `xarray/core/resample.py` as well?
shoyer: This still needs tests, to verify that this argument works.
maaleske: I added tests for groupby. I haven't used resampling so I'm struggling a bit to make a sensible dataset/array for a similar test of the resample method.
spencerkclark: Thanks @maaleske -- following the style of your groupby tests, you could do something like this, e.g. for the DataArray case:
```python
def test_da_resample_func_args():

    def func(arg1, arg2, arg3=0.):
        return arg1.mean('time') + arg2 + arg3

    times = pd.date_range('2000', periods=3, freq='D')
    da = xr.DataArray([1., 1., 1.], coords=[times], dims=['time'])
    expected = xr.DataArray([3., 3., 3.], coords=[times], dims=['time'])
    actual = da.resample(time='D').apply(func, args=(1,), arg3=1)
    assert_identical(actual, expected)
```
maaleske: Thanks @spencerkclark, I had almost the exact same code already written but couldn't figure out the return value `resample.apply` needs from `func` to work. I added the tests to `test_dataarray` and `test_dataset`, since there doesn't appear to be a separate file for `resample` (as there is for `groupby`).
dcherian: lgtm
pep8speaks: Hello @maaleske! Thanks for updating the PR.

 - There are no PEP8 issues in the file [`xarray/core/dataset.py`](https://github.com/maaleske/xarray/blob/8a2da445e16c7d7e7a79d3895fc2c3aa66d4d720/xarray/core/dataset.py) !

 - There are no PEP8 issues in the file [`xarray/core/groupby.py`](https://github.com/maaleske/xarray/blob/8a2da445e16c7d7e7a79d3895fc2c3aa66d4d720/xarray/core/groupby.py) !

 - There are no PEP8 issues in the file [`xarray/core/resample.py`](https://github.com/maaleske/xarray/blob/8a2da445e16c7d7e7a79d3895fc2c3aa66d4d720/xarray/core/resample.py) !

 - There are no PEP8 issues in the file [`xarray/tests/test_dataarray.py`](https://github.com/maaleske/xarray/blob/8a2da445e16c7d7e7a79d3895fc2c3aa66d4d720/xarray/tests/test_dataarray.py) !

 - There are no PEP8 issues in the file [`xarray/tests/test_dataset.py`](https://github.com/maaleske/xarray/blob/8a2da445e16c7d7e7a79d3895fc2c3aa66d4d720/xarray/tests/test_dataset.py) !

 - There are no PEP8 issues in the file [`xarray/tests/test_groupby.py`](https://github.com/maaleske/xarray/blob/8a2da445e16c7d7e7a79d3895fc2c3aa66d4d720/xarray/tests/test_groupby.py) !



Patch:
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
index 6a37e0c0..19e52f3a 100644
--- a/doc/whats-new.rst
+++ b/doc/whats-new.rst
@@ -66,6 +66,10 @@ Enhancements
 - :py:meth:`DataArray.resample` and :py:meth:`Dataset.resample` now supports the
   ``loffset`` kwarg just like Pandas.
   By `Deepak Cherian <https://github.com/dcherian>`_
+- The `apply` methods for `DatasetGroupBy`, `DataArrayGroupBy`,
+  `DatasetResample` and `DataArrayResample` can now pass positional arguments to
+  the applied function.
+  By `Matti Eskelinen <https://github.com/maaleske>`_.
 - 0d slices of ndarrays are now obtained directly through indexing, rather than
   extracting and wrapping a scalar, avoiding unnecessary copying. By `Daniel
   Wennberg <https://github.com/danielwe>`_.
@@ -260,6 +264,9 @@ Announcements of note:
   for more details.
 - We have a new :doc:`roadmap` that outlines our future development plans.
 
+- `Dataset.apply` now properly documents the way `func` is called.
+  By `Matti Eskelinen <https://github.com/maaleske>`_.
+
 Enhancements
 ~~~~~~~~~~~~
 
@@ -267,6 +274,7 @@ Enhancements
   :py:meth:`~xarray.Dataset.differentiate` are newly added.
   (:issue:`1332`)
   By `Keisuke Fujii <https://github.com/fujiisoup>`_.
+
 - Default colormap for sequential and divergent data can now be set via
   :py:func:`~xarray.set_options()`
   (:issue:`2394`)
diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
index 397a0f72..686ffb37 100644
--- a/xarray/core/dataset.py
+++ b/xarray/core/dataset.py
@@ -2953,8 +2953,8 @@ class Dataset(Mapping, ImplementsDatasetReduce, DataWithCoords,
         Parameters
         ----------
         func : function
-            Function which can be called in the form `f(x, **kwargs)` to
-            transform each DataArray `x` in this dataset into another
+            Function which can be called in the form `func(x, *args, **kwargs)`
+            to transform each DataArray `x` in this dataset into another
             DataArray.
         keep_attrs : bool, optional
             If True, the dataset's attributes (`attrs`) will be copied from
diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py
index 58ba4570..9e8c1e6b 100644
--- a/xarray/core/groupby.py
+++ b/xarray/core/groupby.py
@@ -503,7 +503,7 @@ class DataArrayGroupBy(GroupBy, ImplementsArrayReduce):
         new_order = sorted(stacked.dims, key=lookup_order)
         return stacked.transpose(*new_order)
 
-    def apply(self, func, shortcut=False, **kwargs):
+    def apply(self, func, shortcut=False, args=(), **kwargs):
         """Apply a function over each array in the group and concatenate them
         together into a new array.
 
@@ -532,6 +532,8 @@ class DataArrayGroupBy(GroupBy, ImplementsArrayReduce):
             If these conditions are satisfied `shortcut` provides significant
             speedup. This should be the case for many common groupby operations
             (e.g., applying numpy ufuncs).
+        args : tuple, optional
+            Positional arguments passed to `func`.
         **kwargs
             Used to call `func(ar, **kwargs)` for each array `ar`.
 
@@ -544,7 +546,7 @@ class DataArrayGroupBy(GroupBy, ImplementsArrayReduce):
             grouped = self._iter_grouped_shortcut()
         else:
             grouped = self._iter_grouped()
-        applied = (maybe_wrap_array(arr, func(arr, **kwargs))
+        applied = (maybe_wrap_array(arr, func(arr, *args, **kwargs))
                    for arr in grouped)
         return self._combine(applied, shortcut=shortcut)
 
@@ -642,7 +644,7 @@ ops.inject_binary_ops(DataArrayGroupBy)
 
 
 class DatasetGroupBy(GroupBy, ImplementsDatasetReduce):
-    def apply(self, func, **kwargs):
+    def apply(self, func, args=(), **kwargs):
         """Apply a function over each Dataset in the group and concatenate them
         together into a new Dataset.
 
@@ -661,6 +663,8 @@ class DatasetGroupBy(GroupBy, ImplementsDatasetReduce):
         ----------
         func : function
             Callable to apply to each sub-dataset.
+        args : tuple, optional
+            Positional arguments to pass to `func`.
         **kwargs
             Used to call `func(ds, **kwargs)` for each sub-dataset `ar`.
 
@@ -670,7 +674,7 @@ class DatasetGroupBy(GroupBy, ImplementsDatasetReduce):
             The result of splitting, applying and combining this dataset.
         """
         kwargs.pop('shortcut', None)  # ignore shortcut if set (for now)
-        applied = (func(ds, **kwargs) for ds in self._iter_grouped())
+        applied = (func(ds, *args, **kwargs) for ds in self._iter_grouped())
         return self._combine(applied)
 
     def _combine(self, applied):
diff --git a/xarray/core/resample.py b/xarray/core/resample.py
index edf7dfc3..49351efc 100644
--- a/xarray/core/resample.py
+++ b/xarray/core/resample.py
@@ -129,7 +129,7 @@ class DataArrayResample(DataArrayGroupBy, Resample):
                              "('{}')! ".format(self._resample_dim, self._dim))
         super(DataArrayResample, self).__init__(*args, **kwargs)
 
-    def apply(self, func, shortcut=False, **kwargs):
+    def apply(self, func, shortcut=False, args=(), **kwargs):
         """Apply a function over each array in the group and concatenate them
         together into a new array.
 
@@ -158,6 +158,8 @@ class DataArrayResample(DataArrayGroupBy, Resample):
             If these conditions are satisfied `shortcut` provides significant
             speedup. This should be the case for many common groupby operations
             (e.g., applying numpy ufuncs).
+        args : tuple, optional
+            Positional arguments passed on to `func`.
         **kwargs
             Used to call `func(ar, **kwargs)` for each array `ar`.
 
@@ -167,7 +169,7 @@ class DataArrayResample(DataArrayGroupBy, Resample):
             The result of splitting, applying and combining this array.
         """
         combined = super(DataArrayResample, self).apply(
-            func, shortcut=shortcut, **kwargs)
+            func, shortcut=shortcut, args=args, **kwargs)
 
         # If the aggregation function didn't drop the original resampling
         # dimension, then we need to do so before we can rename the proxy
@@ -240,7 +242,7 @@ class DatasetResample(DatasetGroupBy, Resample):
                              "('{}')! ".format(self._resample_dim, self._dim))
         super(DatasetResample, self).__init__(*args, **kwargs)
 
-    def apply(self, func, **kwargs):
+    def apply(self, func, args=(), **kwargs):
         """Apply a function over each Dataset in the groups generated for
         resampling  and concatenate them together into a new Dataset.
 
@@ -259,6 +261,8 @@ class DatasetResample(DatasetGroupBy, Resample):
         ----------
         func : function
             Callable to apply to each sub-dataset.
+        args : tuple, optional
+            Positional arguments passed on to `func`.
         **kwargs
             Used to call `func(ds, **kwargs)` for each sub-dataset `ar`.
 
@@ -268,7 +272,7 @@ class DatasetResample(DatasetGroupBy, Resample):
             The result of splitting, applying and combining this dataset.
         """
         kwargs.pop('shortcut', None)  # ignore shortcut if set (for now)
-        applied = (func(ds, **kwargs) for ds in self._iter_grouped())
+        applied = (func(ds, *args, **kwargs) for ds in self._iter_grouped())
         combined = self._combine(applied)
 
         return combined.rename({self._resample_dim: self._dim})


Test patch:
diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
index 0672861d..fe05da16 100644
--- a/xarray/tests/test_dataarray.py
+++ b/xarray/tests/test_dataarray.py
@@ -2295,6 +2295,17 @@ class TestDataArray(object):
         with raises_regex(ValueError, 'index must be monotonic'):
             array[[2, 0, 1]].resample(time='1D')
 
+    def test_da_resample_func_args(self):
+
+        def func(arg1, arg2, arg3=0.):
+            return arg1.mean('time') + arg2 + arg3
+
+        times = pd.date_range('2000', periods=3, freq='D')
+        da = xr.DataArray([1., 1., 1.], coords=[times], dims=['time'])
+        expected = xr.DataArray([3., 3., 3.], coords=[times], dims=['time'])
+        actual = da.resample(time='D').apply(func, args=(1.,), arg3=1.)
+        assert_identical(actual, expected)
+
     @requires_cftime
     def test_resample_cftimeindex(self):
         cftime = _import_cftime()
diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
index d4253ae4..521f2395 100644
--- a/xarray/tests/test_dataset.py
+++ b/xarray/tests/test_dataset.py
@@ -2886,6 +2886,19 @@ class TestDataset(object):
         with raises_regex(TypeError, r'resample\(\) no longer supports'):
             ds.resample('1D', dim='time')
 
+    def test_ds_resample_apply_func_args(self):
+
+        def func(arg1, arg2, arg3=0.):
+            return arg1.mean('time') + arg2 + arg3
+
+        times = pd.date_range('2000', freq='D', periods=3)
+        ds = xr.Dataset({'foo': ('time', [1., 1., 1.]),
+                         'time': times})
+        expected = xr.Dataset({'foo': ('time', [3., 3., 3.]),
+                               'time': times})
+        actual = ds.resample(time='D').apply(func, args=(1.,), arg3=1.)
+        assert_identical(expected, actual)
+
     def test_to_array(self):
         ds = Dataset(OrderedDict([('a', 1), ('b', ('x', [1, 2, 3]))]),
                      coords={'c': 42}, attrs={'Conventions': 'None'})
diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py
index 8ace55be..205bc931 100644
--- a/xarray/tests/test_groupby.py
+++ b/xarray/tests/test_groupby.py
@@ -85,4 +85,26 @@ def test_groupby_input_mutation():
     assert_identical(array, array_copy)  # should not modify inputs
 
 
+def test_da_groupby_apply_func_args():
+
+    def func(arg1, arg2, arg3=0):
+        return arg1 + arg2 + arg3
+
+    array = xr.DataArray([1, 1, 1], [('x', [1, 2, 3])])
+    expected = xr.DataArray([3, 3, 3], [('x', [1, 2, 3])])
+    actual = array.groupby('x').apply(func, args=(1,), arg3=1)
+    assert_identical(expected, actual)
+
+
+def test_ds_groupby_apply_func_args():
+
+    def func(arg1, arg2, arg3=0):
+        return arg1 + arg2 + arg3
+
+    dataset = xr.Dataset({'foo': ('x', [1, 1, 1])}, {'x': [1, 2, 3]})
+    expected = xr.Dataset({'foo': ('x', [3, 3, 3])}, {'x': [1, 2, 3]})
+    actual = dataset.groupby('x').apply(func, args=(1,), arg3=1)
+    assert_identical(expected, actual)
+
+
 # TODO: move other groupby tests from test_dataset and test_dataarray over here