import numpy_groupies as npg
from numpy import array, nan
a = array([[5.0, 6.0, 7.0, 8.0, 9.0], [0.0, 0.0, 0.0, nan, 0.0]])
#mask = np.isnan(a)
#a[mask] = 0
result = npg.aggregate_numpy.aggregate(
array([1, 1, 1, 0, 0]),
a,
func="cumsum",
axis=-1,
)
#result[mask] = np.nan
result
array([[ 5., 11., 18., 8., 17.],
[nan, nan, nan, nan, nan]])
We just need to do the nan replacement as with nancumsum
We just need to do the nan replacement as with
nancumsum