numpy - Reset cumsum if over limit (python) -

numpy - Reset cumsum if over limit (python) -

the next numpy snippet homecoming cumsum of input array, resets every time nan encountered.

v = np.array([1., 1., 1., np.nan, 1., 1., 1., 1., np.nan, 1.]) n = np.isnan(v) = ~n c = np.cumsum(a) d = np.diff(np.concatenate(([0.], c[n]))) v[n] = -d result = np.cumsum(v)

in similar fashion, how can calculate cumsum resets if cumsum on value using vectorized pandas or numpy operations?

e.g. limit = 5, in = [1,1,1,1,1,1,1,1,1,1], out = [1,2,3,4,5,1,2,3,4,5]

if numbers in array positive, simplest utilize cumsum() , modulo operator:

>>> = np.array([1,1,1,1,1,1,1,1,1,1]) >>> limit = 5 >>> x = a.cumsum() % limit >>> x array([1, 2, 3, 4, 0, 1, 2, 3, 4, 0])

you can set 0 values limit desired array:

>>> x[x == 0] = limit >>> x array([1, 2, 3, 4, 5, 1, 2, 3, 4, 5])

here's 1 possible general solution using pandas' expanding_apply method. (i've not tested extensively...)

first define modified cumsum function:

import pandas pd def cumsum_limit(x): q = np.sum(x[:-1]) if q > 0: q = q%5 r = x[-1] if q+r <= 5: homecoming q+r elif (q+r)%5 == 0: homecoming 5 else: homecoming (q+r)%5 = np.array([1,1,1,1,1,1,1,1,1,1]) # illustration array

apply function array this:

>>> pd.expanding_apply(a, lambda x: cumsum_limit(x)) array([ 1., 2., 3., 4., 5., 1., 2., 3., 4., 5.])

here's function applied more interesting series:

>>> s = pd.series([3, -8, 4, 5, -3, 501, 7, -100, 98, 3]) >>> pd.expanding_apply(s, lambda x: cumsum_limit(x)) 0 3 1 -5 2 -1 3 4 4 1 5 2 6 4 7 -96 8 2 9 5 dtype: float64

python numpy pandas