Numpy cumsum considering NaNs

arrays, cumsum, nan, numpy, python

Solution

How about (for not-too-big arrays):

In [34]: import numpy as np

In [35]: a = np.array([1,4,1,np.nan,2,np.nan])

In [36]: a*0 + np.nan_to_num(a).cumsum()
Out[36]: array([  1.,   5.,   6.,  nan,   8.,  nan])

Problem

I am looking for a succinct way to go from: ``` a = numpy.array([1,4,1,numpy.nan,2,numpy.nan]) ``` to: ``` b = numpy.array([1,5,6,numpy.nan,8,numpy.nan]) ``` The best I can do currently is: ``` b = numpy.insert(numpy.cumsum(a[numpy.isfinite(a)]), (numpy.argwhere(numpy.isnan(a)) - numpy.arange(len(numpy.argwhere(numpy.isnan(a))))), numpy.nan) ``` Is there a shorter way to accomplish the same? What about doing a cumsum along an axis of a 2D array?

Original source