8000 Adds nancumsum, nancumprod to xarray functions · pydata/xarray@9242854 · GitHub
[go: up one dir, main page]
Skip to content 10000

Commit 9242854

Browse files
pwolframpwolfram
committed
Adds nancumsum, nancumprod to xarray functions
1 parent 6174efb commit 9242854

File tree

3 files changed

+166
-128
lines changed

3 files changed

+166
-128
lines changed

‎xarray/core/npcompat.py

Lines changed: 125 additions & 125 deletions
Original file line numberDiff line numberDiff line change
@@ -253,128 +253,128 @@ def nanprod(a, axis=None, dtype=None, out=None, keepdims=0):
253253
return np.prod(a, axis=axis, dtype=dtype, out=out, keepdims=keepdims)
254254

255255

256-
def nancumsum(a, axis=None, dtype=None, out=None):
257-
"""
258-
Return the cumulative sum of array elements over a given axis treating Not a
259-
Numbers (NaNs) as zero. The cumulative sum does not change when NaNs are
260-
encountered and leading NaNs are replaced by zeros.
261-
262-
Zeros are returned for slices that are all-NaN or empty.
263-
264-
.. versionadded:: 1.12.0
265-
266-
Parameters
267-
----------
268-
a : array_like
269-
Input array.
270-
axis : int, optional
271-
Axis along which the cumulative sum is computed. The default
272-
(None) is to compute the cumsum over the flattened array.
273-
dtype : dtype, optional
274-
Type of the returned array and of the accumulator in which the
275-
elements are summed. If `dtype` is not specified, it defaults
276-
to the dtype of `a`, unless `a` has an integer dtype with a
277-
precision less than that of the default platform integer. In
278-
that case, the default platform integer is used.
279-
out : ndarray, optional
280-
Alternative output array in which to place the result. It must
281-
have the same shape and buffer length as the expected output
282-
but the type will be cast if necessary. See `doc.ufuncs`
283-
(Section "Output arguments") for more details.
284-
285-
Returns
286-
-------
287-
nancumsum : ndarray.
288-
A new array holding the result is returned unless `out` is
289-
specified, in which it is returned. The result has the same
290-
size as `a`, and the same shape as `a` if `axis` is not None
291-
or `a` is a 1-d array.
292-
293-
See Also
294-
--------
295-
numpy.cumsum : Cumulative sum across array propagating NaNs.
296-
isnan : Show which elements are NaN.
297-
298-
Examples
299-
--------
300-
>>> np.nancumsum(1)
301-
array([1])
302-
>>> np.nancumsum([1])
303-
array([1])
304-
>>> np.nancumsum([1, np.nan])
305-
array([ 1., 1.])
306-
>>> a = np.array([[1, 2], [3, np.nan]])
307-
>>> np.nancumsum(a)
308-
array([ 1., 3., 6., 6.])
309-
>>> np.nancumsum(a, axis=0)
310-
array([[ 1., 2.],
311-
[ 4., 2.]])
312-
>>> np.nancumsum(a, axis=1)
313-
array([[ 1., 3.],
314-
[ 3., 3.]])
315-
316-
"""
317-
a, mask = _replace_nan(a, 0)
318-
return np.cumsum(a, axis=axis, dtype=dtype, out=out)
319-
320-
321-
def nancumprod(a, axis=None, dtype=None, out=None):
322-
"""
323-
Return the cumulative product of array elements over a given axis treating Not a
324-
Numbers (NaNs) as one. The cumulative product does not change when NaNs are
325-
encountered and leading NaNs are replaced by ones.
326-
327-
Ones are returned for slices that are all-NaN or empty.
328-
329-
.. versionadded:: 1.12.0
330-
331-
Parameters
332-
----------
333-
a : array_like
334-
Input array.
335-
axis : int, optional
336-
Axis along which the cumulative product is computed. By default
337-
the input is flattened.
338-
dtype : dtype, optional
339-
Type of the returned array, as well as of the accumulator in which
340-
the elements are multiplied. If *dtype* is not specified, it
341-
defaults to the dtype of `a`, unless `a` has an integer dtype with
342-
a precision less than that of the default platform integer. In
343-
that case, the default platform integer is used instead.
344-
out : ndarray, optional
345-
Alternative output array in which to place the result. It must
346-
have the same shape and buffer length as the expected output
347-
but the type of the resulting values will be cast if necessary.
348-
349-
Returns
350-
-------
351-
nancumprod : ndarray
352-
A new array holding the result is returned unless `out` is
353-
specified, in which case it is returned.
354-
355-
See Also
356-
--------
357-
numpy.cumprod : Cumulative product across array propagating NaNs.
358-
isnan : Show which elements are NaN.
359-
360-
Examples
361-
--------
362-
>>> np.nancumprod(1)
363-
array([1])
364-
>>> np.nancumprod([1])
365-
array([1])
366-
>>> np.nancumprod([1, np.nan])
367-
array([ 1., 1.])
368-
>>> a = np.array([[1, 2], [3, np.nan]])
369-
>>> np.nancumprod(a)
370-
array([ 1., 2., 6., 6.])
371-
>>> np.nancumprod(a, axis=0)
372-
array([[ 1., 2.],
373-
[ 3., 2.]])
374-
>>> np.nancumprod(a, axis=1)
375-
array([[ 1., 2.],
376-
[ 3., 3.]])
377-
378-
"""
379-
a, mask = _replace_nan(a, 1)
380-
return np.cumprod(a, axis=axis, dtype=dtype, out=out)
256+
def nancumsum(a, axis=None, dtype=None, out=None):
257+
"""
258+
Return the cumulative sum of array elements over a given axis treating Not a
259+
Numbers (NaNs) as zero. The cumulative sum does not change when NaNs are
260+
encountered and leading NaNs are replaced by zeros.
261+
262+
Zeros are returned for slices that are all-NaN or empty.
263+
264+
.. versionadded:: 1.12.0
265+
266+
Parameters
267+
----------
268+
a : array_like
269+
Input array.
270+
axis : int, optional
271+
Axis along which the cumulative sum is computed. The default
272+
(None) is to compute the cumsum over the flattened array.
273+
dtype : dtype, optional
274+
Type of the returned array and of the accumulator in which the
275+
elements are summed. If `dtype` is not specified, it defaults
276+
to the dtype of `a`, unless `a` has an integer dtype with a
277+
precision less than that of the default platform integer. In
278+
that case, the default platform integer is used.
279+
out : ndarray, optional
280+
Alternative output array in which to place the result. It must
281+
have the same shape and buffer length as the expected output
282+
but the type will be cast if necessary. See `doc.ufuncs`
283+
(Section "Output arguments") for more details.
284+
285+
Returns
286+
-------
287+
nancumsum : ndarray.
288+
A new array holding the result is returned unless `out` is
289+
specified, in which it is returned. The result has the same
290+
size as `a`, and the same shape as `a` if `axis` is not None
291+
or `a` is a 1-d array.
292+
293+
See Also
294+
--------
295+
numpy.cumsum : Cumulative sum across array propagating NaNs.
296+
isnan : Show which elements are NaN.
297+
298+
Examples
299+
--------
300+
>>> np.nancumsum(1)
301+
array([1])
302+
>>> np.nancumsum([1])
303+
array([1])
304+
>>> np.nancumsum([1, np.nan])
305+
array([ 1., 1.])
306+
>>> a = np.array([[1, 2], [3, np.nan]])
307+
>>> np.nancumsum(a)
308+
array([ 1., 3., 6., 6.])
309+
>>> np.nancumsum(a, axis=0)
310+
array([[ 1., 2.],
311+
[ 4., 2.]])
312+
>>> np.nancumsum(a, axis=1)
313+
array([[ 1., 3.],
314+
[ 3., 3.]])
315+
316+
"""
317+
a, mask = _replace_nan(a, 0)
318+
return np.cumsum(a, axis=axis, dtype=dtype, out=out)
319+
320+
321+
def nancumprod(a, axis=None, dtype=None, out=None):
322+
"""
323+
Return the cumulative product of array elements over a given axis treating Not a
324+
Numbers (NaNs) as one. The cumulative product does not change when NaNs are
325+
encountered and leading NaNs are replaced by ones.
326+
327+
Ones are returned for slices that are all-NaN or empty.
328+
329+
.. versionadded:: 1.12.0
330+
331+
Parameters
332+
----------
333+
a : array_like
334+
Input array.
335+
axis : int, optional
336+
Axis along which the cumulative product is computed. By default
337+
the input is flattened.
338+
dtype : dtype, optional
339+
Type of the returned array, as well as of the accumulator in which
340+
the elements are multiplied. If *dtype* is not specified, it
341+
defaults to the dtype of `a`, unless `a` has an integer dtype with
342+
a precision less than that of the default platform integer. In
343+
that case, the default platform integer is used instead.
344+
out : ndarray, optional
345+
Alternative output array in which to place the result. It must
346+
have the same shape and buffer length as the expected output
347+
but the type of the resulting values will be cast if necessary.
348+
349+
Returns
350+
-------
351+
nancumprod : ndarray
352+
A new array holding the result is returned unless `out` is
353+
specified, in which case it is returned.
354+
355+
See Also
356+
--------
357+
numpy.cumprod : Cumulative product across array propagating NaNs.
358+
isnan : Show which elements are NaN.
359+
360+
Examples
361+
--------
362+
>>> np.nancumprod(1)
363+
array([1])
364+
>>> np.nancumprod([1])
365+
array([1])
366+
>>> np.nancumprod([1, np.nan])
367+
array([ 1., 1.])
368+
>>> a = np.array([[1, 2], [3, np.nan]])
369+
>>> np.nancumprod(a)
370+
array([ 1., 2., 6., 6.])
371+
>>> np.nancumprod(a, axis=0)
372+
array([[ 1., 2.],
373+
[ 3., 2.]])
374+
>>> np.nancumprod(a, axis=1)
375+
array([[ 1., 2.],
376+
[ 3., 3.]])
377+
378+
"""
379+
a, mask = _replace_nan(a, 1)
380+
return np.cumprod(a, axis=axis, dtype=dtype, out=out)

‎xarray/core/ops.py

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,11 +44,11 @@
4444
# methods which remove an axis
4545
REDUCE_METHODS = ['all', 'any']
4646
NAN_REDUCE_METHODS = ['argmax', 'argmin', 'max', 'min', 'mean', 'prod', 'sum',
47-
'std', 'var', 'median']
47+
'std', 'var', 'median', 'cumsum', 'cumprod']
4848
BOTTLENECK_ROLLING_METHODS = {'move_sum': 'sum', 'move_mean': 'mean',
4949
'move_std': 'std', 'move_min': 'min',
5050
'move_max': 'max'}
51-
# TODO: wrap cumprod/cumsum, take, dot, sort
51+
# TODO: wrap take, dot, sort
5252

5353

5454
def _dask_or_eager_func(name, eager_module=np, list_of_args=False,
@@ -336,6 +336,8 @@ def f(values, axis=None, skipna=None, **kwargs):
336336

337337

338338
_prod = _dask_or_eager_func('prod')
339+
_cumsum = _dask_or_eager_func('cumsum')
340+
_cumprod = _dask_or_eager_func('cumprod')
339341

340342

341343
def prod(values, axis=None, skipna=None, **kwargs):
@@ -349,6 +351,28 @@ def prod(values, axis=None, skipna=None, **kwargs):
349351
return _prod(values, axis=axis, **kwargs)
350352
prod.numeric_only = True
351353

354+
def cumsum(values, axis=None, skipna=None, **kwargs):
355+
if skipna or (skipna is None and values.dtype.kind == 'f'):
356+
if values.dtype.kind not in ['i', 'f']:
357+
raise NotImplementedError(
358+
'skipna=True not yet implemented for cumsum with dtype %s'
359+
% values.dtype)
360+
_fail_on_dask_array_input_skipna(values)
361+
return npcompat.nancumsum(values, axis=axis, **kwargs)
362+
return _cumsum(values, axis=axis, **kwargs)
363+
cumsum.numeric_only = True
364+
365+
def cumprod(values, axis=None, skipna=None, **kwargs):
366+
if skipna or (skipna is None and values.dtype.kind == 'f'):
367+
if values.dtype.kind not in ['i', 'f']:
368+
raise NotImplementedError(
369+
'skipna=True not yet implemented for cumprod with dtype %s'
370+
% values.dtype)
371+
_fail_on_dask_array_input_skipna(values)
372+
return npcompat.nancumprod(values, axis=axis, **kwargs)
373+
return _cumprod(values, axis=axis, **kwargs)
374+
cumprod.numeric_only = True
375+
352376

353377
def first(values, axis, skipna=None):
354378
"""Return the first non-NA elements in this array along the given axis

‎xarray/core/variable.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -900,9 +900,23 @@ def reduce(self, func, dim=None, axis=None, keep_attrs=False,
900900

901901
removed_axes = (range(self.ndim) if axis is None
902902
else np.atleast_1d(axis) % self.ndim)
903-
dims = [dim for n, dim in enumerate(self.dims)
903+
dims = [adim for n, adim in enumerate(self.dims)
904904
if n not in removed_axes]
905905

906+
if 'cum' in func.__name__:
907+
if dim is None:
908+
raise ValueError("must supply a 'dim' argument to %s"%(func.__name__))
909+
910+
def safe_shape(val):
911+
if type(val) is np.ndarray:
912+
return val.shape
913+
else:
914+
# single value
915+
return ()
916+
917+
if safe_shape(data) == safe_shape(self.data):
918+
dims = self.dims
919+
906920
attrs = self._attrs if keep_attrs else None
907921

908922
return Variable(dims, data, attrs=attrs)

0 commit comments

Comments
 (0)
0