From 3bd704a4815ad2281e61eedcee3c7935789d410b Mon Sep 17 00:00:00 2001 From: Keisuke Fujii Date: Sun, 21 Jan 2018 03:29:01 +0900 Subject: [PATCH] Add dtype support for reduce methods. (#1841) * Add dtype support for reduce methods. * Use bottleneck if dtype == values.dtype * Change get -> pop * Update comments. * float128 -> float16 * ValueError -> TypeError * update test --- doc/whats-new.rst | 4 +++- xarray/core/duck_array_ops.py | 12 +++++++----- xarray/tests/test_dataarray.py | 18 ++++++++++++++++++ 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index ab3d9733856..cf0f3fc876f 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -27,9 +27,11 @@ Documentation Enhancements ~~~~~~~~~~~~ +- reduce methods such as :py:func:`DataArray.sum()` now accepts ``dtype`` + arguments. (:issue:`1838`) + By `Keisuke Fujii `_. - Added nodatavals attribute to DataArray when using :py:func:`~xarray.open_rasterio`. (:issue:`1736`). By `Alan Snow `_. - - :py:func:`~plot.contourf()` learned to contour 2D variables that have both a 1D co-ordinate (e.g. time) and a 2D co-ordinate (e.g. depth as a function of time) (:issue:`1737`). diff --git a/xarray/core/duck_array_ops.py b/xarray/core/duck_array_ops.py index 2f97de4a1ba..2058ce86a99 100644 --- a/xarray/core/duck_array_ops.py +++ b/xarray/core/duck_array_ops.py @@ -175,11 +175,11 @@ def _create_nan_agg_method(name, numeric_only=False, np_compat=False, no_bottleneck=False, coerce_strings=False, keep_dims=False): def f(values, axis=None, skipna=None, **kwargs): - # ignore keyword args inserted by np.mean and other numpy aggregators - # automatically: - kwargs.pop('dtype', None) - kwargs.pop('out', None) + if kwargs.pop('out', None) is not None: + raise TypeError('`out` is not valid for {}'.format(name)) + # If dtype is supplied, we use numpy's method. + dtype = kwargs.get('dtype', None) values = asarray(values) if coerce_strings and values.dtype.kind in 'SU': @@ -192,7 +192,8 @@ def f(values, axis=None, skipna=None, **kwargs): % (name, values.dtype)) nanname = 'nan' + name if (isinstance(axis, tuple) or not values.dtype.isnative or - no_bottleneck): + no_bottleneck or + (dtype is not None and np.dtype(dtype) != values.dtype)): # bottleneck can't handle multiple axis arguments or non-native # endianness if np_compat: @@ -200,6 +201,7 @@ def f(values, axis=None, skipna=None, **kwargs): else: eager_module = np else: + kwargs.pop('dtype', None) eager_module = bn func = _dask_or_eager_func(nanname, eager_module) using_numpy_nan_func = (eager_module is np or diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 5524b4c971d..39b3109c295 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -1784,6 +1784,24 @@ def test_reduce(self): dims=['x', 'y']).mean('x') assert_equal(actual, expected) + def test_reduce_dtype(self): + coords = {'x': [-1, -2], 'y': ['ab', 'cd', 'ef'], + 'lat': (['x', 'y'], [[1, 2, 3], [-1, -2, -3]]), + 'c': -999} + orig = DataArray([[-1, 0, 1], [-3, 0, 3]], coords, dims=['x', 'y']) + + for dtype in [np.float16, np.float32, np.float64]: + assert orig.astype(float).mean(dtype=dtype).dtype == dtype + + def test_reduce_out(self): + coords = {'x': [-1, -2], 'y': ['ab', 'cd', 'ef'], + 'lat': (['x', 'y'], [[1, 2, 3], [-1, -2, -3]]), + 'c': -999} + orig = DataArray([[-1, 0, 1], [-3, 0, 3]], coords, dims=['x', 'y']) + + with pytest.raises(TypeError): + orig.mean(out=np.ones(orig.shape)) + # skip due to bug in older versions of numpy.nanpercentile def test_quantile(self): for q in [0.25, [0.50], [0.25, 0.75]]: