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

Commit ce83d37

Browse files
pwolframpwolfram
committed
Adds nancumsum, nancumprod to xarray functions
1 parent 74ece06 commit ce83d37

File tree

2 files changed

+160
-126
lines changed

2 files changed

+160
-126
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: 35 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,14 +41,15 @@
4141
# wrapped in an Variable/DataArray
4242
NUMPY_UNARY_METHODS = ['astype', 'argsort', 'clip', 'conj', 'conjugate']
4343
PANDAS_UNARY_FUNCTIONS = ['isnull', 'notnull']
44+
NAN_METHODS = ['cumsum','cumprod']
4445
# methods which remove an axis
4546
REDUCE_METHODS = ['all', 'any']
4647
NAN_REDUCE_METHODS = ['argmax', 'argmin', 'max', 'min', 'mean', 'prod', 'sum',
4748
'std', 'var', 'median']
4849
BOTTLENECK_ROLLING_METHODS = {'move_sum': 'sum', 'move_mean': 'mean',
4950
'move_std': 'std', 'move_min': 'min',
5051
'move_max': 'max'}
51-
# TODO: wrap cumprod/cumsum, take, dot, sort
52+
# TODO: wrap take, dot, sort
5253

5354

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

337338

338339
_prod = _dask_or_eager_func('prod')
340+
_cumsum = _dask_or_eager_func('cumsum')
341+
_cumprod = _dask_or_eager_func('cumprod')
339342

340343

341344
def prod(values, axis=None, skipna=None, **kwargs):
@@ -349,6 +352,28 @@ def prod(values, axis=None, skipna=None, **kwargs):
349352
return _prod(values, axis=axis, **kwargs)
350353
prod.numeric_only = True
351354

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

353378
def first(values, axis, skipna=None):
354379
"""Return the first non-NA elements in this array along the given axis
@@ -384,6 +409,14 @@ def inject_reduce_methods(cls):
384409
extra_args=cls._reduce_extra_args_docstring)
385410
setattr(cls, name, func)
386411

412+
def inject_nan_methods(cls):
413+
methods = ([(name, globals()[name]) for name in NAN_METHODS])
414+
for name, f in methods:
415+
func = f
416+
func.__name__ = name
417+
func.__doc__ = getattr(np.ndarray, name).__doc__
418+
setattr(cls, name, func)
419+
387420

388421
def op_str(name):
389422
return '__%s__' % name
@@ -454,6 +487,7 @@ def inject_all_ops_and_reduce_methods(cls, priority=50, array_only=True):
454487
setattr(cls, name, _values_method_wrapper(name))
455488

456489
inject_reduce_methods(cls)
490+
inject_nan_methods(cls)
457491

458492

459493
def inject_bottleneck_rolling_methods(cls):

0 commit comments

Comments
 (0)
0