¿Puedo realizar una acumulación dinámica de filas en pandas?

Si tengo el siguiente dataframe, derivado así: df = pd.DataFrame(np.random.randint(0, 10, size=(10, 1)))

  0 0 0 1 2 2 8 3 1 4 0 5 0 6 7 7 0 8 2 9 2 

¿Existe una forma eficiente de cumsum filas con un límite y cada vez que se alcanza este límite, para iniciar una nueva cumsum ? Una vez que se alcanza cada límite (sin embargo, muchas filas), se crea una fila con la sum total.

A continuación, he creado un ejemplo de una función que hace esto, pero es muy lenta, especialmente cuando el dataframe es muy grande. No me gusta que mi función esté en bucle y busco una manera de hacerlo más rápido (supongo que es una forma sin bucle).

 def foo(df, max_value): last_value = 0 storage = [] for index, row in df.iterrows(): this_value = np.nansum([row[0], last_value]) if this_value >= max_value: storage.append((index, this_value)) this_value = 0 last_value = this_value return storage 

Si modificas mi función así: foo(df, 5) En el contexto anterior, devuelve:

  0 2 10 6 8 

El bucle no se puede evitar, pero se puede numba utilizando el numba de njit :

 from numba import njit, prange @njit def dynamic_cumsum(seq, index, max_value): cumsum = [] running = 0 for i in prange(len(seq)): if running > max_value: cumsum.append([index[i], running]) running = 0 running += seq[i] cumsum.append([index[-1], running]) return cumsum 

El índice se requiere aquí, asumiendo que su índice no es numérico / monótonamente creciente.

 %timeit foo(df, 5) 1.24 ms ± 41.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) %timeit dynamic_cumsum(df.iloc(axis=1)[0].values, df.index.values, 5) 77.2 µs ± 4.01 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) 

Si el índice es de tipo Int64Index , puede Int64Index a:

 @njit def dynamic_cumsum2(seq, max_value): cumsum = [] running = 0 for i in prange(len(seq)): if running > max_value: cumsum.append([i, running]) running = 0 running += seq[i] cumsum.append([i, running]) return cumsum lst = dynamic_cumsum2(df.iloc(axis=1)[0].values, 5) pd.DataFrame(lst, columns=['A', 'B']).set_index('A') B A 3 10 7 8 9 4 

 %timeit foo(df, 5) 1.23 ms ± 30.6 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each) %timeit dynamic_cumsum2(df.iloc(axis=1)[0].values, 5) 71.4 µs ± 1.4 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) 

njit Functions Performance

 perfplot.show( setup=lambda n: pd.DataFrame(np.random.randint(0, 10, size=(n, 1))), kernels=[ lambda df: list(cumsum_limit_nb(df.iloc[:, 0].values, 5)), lambda df: dynamic_cumsum2(df.iloc[:, 0].values, 5) ], labels=['cumsum_limit_nb', 'dynamic_cumsum2'], n_range=[2**k for k in range(0, 17)], xlabel='N', logx=True, logy=True, equality_check=None # TODO - update when @jpp adds in the final `yield` ) 

El gráfico log-log muestra que la función del generador es más rápida para entradas más grandes:

introduzca la descripción de la imagen aquí

Una posible explicación es que, a medida que N aumenta, la sobrecarga de agregar a una lista creciente en dynamic_cumsum2 vuelve prominente. Mientras cumsum_limit_nb solo tiene que yield .

Un bucle no es necesariamente malo. El truco es asegurarse de que se realiza en objetos de bajo nivel. En este caso, puedes usar Numba o Cython. Por ejemplo, usando un generador con numba.njit :

 from numba import njit @njit def cumsum_limit(A, limit=5): count = 0 for i in range(A.shape[0]): count += A[i] if count > limit: yield i, count count = 0 idx, vals = zip(*cumsum_limit(df[0].values)) res = pd.Series(vals, index=idx) 

Para demostrar los beneficios de rendimiento de la comstackción JIT con Numba:

 import pandas as pd, numpy as np from numba import njit df = pd.DataFrame({0: [0, 2, 8, 1, 0, 0, 7, 0, 2, 2]}) @njit def cumsum_limit_nb(A, limit=5): count = 0 for i in range(A.shape[0]): count += A[i] if count > limit: yield i, count count = 0 def cumsum_limit(A, limit=5): count = 0 for i in range(A.shape[0]): count += A[i] if count > limit: yield i, count count = 0 n = 10**4 df = pd.concat([df]*n, ignore_index=True) %timeit list(cumsum_limit_nb(df[0].values)) # 4.19 ms ± 90.4 µs per loop %timeit list(cumsum_limit(df[0].values)) # 58.3 ms ± 194 µs per loop 

enfoque más simple:

 def dynamic_cumsum(seq,limit): res=[] cs=seq.cumsum() for i, e in enumerate(cs): if cs[i] >limit: res.append([i,e]) cs[i+1:] -= e if res[-1][0]==i: return res res.append([i,e]) return res 

resultado:

 x=dynamic_cumsum(df[0].values,5) x >>[[2, 10], [6, 8], [9, 4]]