Solución:
Una observación que podría resultar útil es que no es necesario ordenar todos los valores en cada paso. Por el contrario, si se asegura de que la ventana siempre esté ordenada, todo lo que necesita hacer es insertar el nuevo valor en el lugar relevante y eliminar el anterior de donde estaba, las cuales son operaciones que se pueden realizar en O (log_2 (tamaño_ventana)) usando bisect. En la práctica, esto se parecería a
def rolling_mean(data):
x = sorted(data[:49])
res = np.repeat(np.nan, len(data))
for i in range(49, len(data)):
if i != 49:
del x[bisect.bisect_left(x, data[i - 50])]
bisect.insort_right(x, data[i])
res[i] = np.mean(x[3:47])
return res
Ahora, el beneficio adicional en este caso resulta ser menor que lo que se gana con la vectorización que scipy.stats.trim_mean confía, y en particular, esto seguirá siendo más lento que la solución de @ ChrisA, pero es un punto de partida útil para una mayor optimización del rendimiento.
> data = pd.Series(np.random.randint(0, 1000, 50000))
> %timeit data.rolling(50).apply(lambda w: trim_mean(w, 0.06))
727 ms ± 34.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
> %timeit rolling_mean(data.values)
812 ms ± 42.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
En particular, el jitter de Numba, que a menudo es útil en situaciones como estas, tampoco proporciona ningún beneficio:
> from numba import jit
> rolling_mean_jit = jit(rolling_mean)
> %timeit rolling_mean_jit(data.values)
1.05 s ± 183 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
El siguiente enfoque, aparentemente lejos de ser óptimo, supera a los otros dos enfoques considerados anteriormente:
def rolling_mean_np(data):
res = np.repeat(np.nan, len(data))
for i in range(len(data)-49):
x = np.sort(data[i:i+50])
res[i+49] = x[3:47].mean()
return res
Momento:
> %timeit rolling_mean_np(data.values)
564 ms ± 4.44 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Además, esta vez, compilación JIT lo hace ayuda:
> rolling_mean_np_jit = jit(rolling_mean_np)
> %timeit rolling_mean_np_jit(data.values)
94.9 ms ± 605 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
Ya que estamos en eso, verifiquemos rápidamente que esto realmente hace lo que esperamos que haga:
> np.all(rolling_mean_np_jit(data.values)[49:] == data.rolling(50).apply(lambda w: trim_mean(w, 0.06)).values[49:])
True
De hecho, al ayudar un poco al clasificador, podemos exprimir otro factor de 2, reduciendo el tiempo total a 57 ms:
def rolling_mean_np_manual(data):
x = np.sort(data[:50])
res = np.repeat(np.nan, len(data))
for i in range(50, len(data)+1):
res[i-1] = x[3:47].mean()
if i != len(data):
idx_old = np.searchsorted(x, data[i-50])
x[idx_old] = data[i]
x.sort()
return res
> %timeit rolling_mean_np_manual(data.values)
580 ms ± 23 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
> rolling_mean_np_manual_jit = jit(rolling_mean_np_manual)
> %timeit rolling_mean_np_manual_jit(data.values)
57 ms ± 5.89 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
> np.all(rolling_mean_np_manual_jit(data.values)[49:] == data.rolling(50).apply(lambda w: trim_mean(w, 0.06)).values[49:])
True
Ahora, la “clasificación” que se lleva a cabo en este ejemplo, por supuesto, se reduce a colocar el nuevo elemento en el lugar correcto, mientras se cambia todo lo que hay en el medio en uno. Hacer esto a mano hará que el código Python puro sea más lento, pero la versión jitted gana otro factor de 2, llevándonos por debajo de 30 ms:
def rolling_mean_np_shift(data):
x = np.sort(data[:50])
res = np.repeat(np.nan, len(data))
for i in range(50, len(data)+1):
res[i-1] = x[3:47].mean()
if i != len(data):
idx_old, idx_new = np.searchsorted(x, [data[i-50], data[i]])
if idx_old < idx_new:
x[idx_old:idx_new-1] = x[idx_old+1:idx_new]
x[idx_new-1] = data[i]
elif idx_new < idx_old:
x[idx_new+1:idx_old+1] = x[idx_new:idx_old]
x[idx_new] = data[i]
else:
x[idx_new] = data[i]
return res
> %timeit rolling_mean_np_shift(data.values)
937 ms ± 97.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
> rolling_mean_np_shift_jit = jit(rolling_mean_np_shift)
> %timeit rolling_mean_np_shift_jit(data.values)
26.4 ms ± 693 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
> np.all(rolling_mean_np_shift_jit(data.values)[49:] == data.rolling(50).apply(lambda w: trim_mean(w, 0.06)).values[49:])
True
En este punto, la mayor parte del tiempo se dedica a np.searchsorted, así que hagamos que la búsqueda sea compatible con JIT. Adoptando el código fuente para bisect, dejamos
@jit
def binary_search(a, x):
lo = 0
hi = 50
while lo < hi:
mid = (lo+hi)//2
if a[mid] < x: lo = mid+1
else: hi = mid
return lo
@jit
def rolling_mean_np_jitted_search(data):
x = np.sort(data[:50])
res = np.repeat(np.nan, len(data))
for i in range(50, len(data)+1):
res[i-1] = x[3:47].mean()
if i != len(data):
idx_old = binary_search(x, data[i-50])
idx_new = binary_search(x, data[i])
if idx_old < idx_new:
x[idx_old:idx_new-1] = x[idx_old+1:idx_new]
x[idx_new-1] = data[i]
elif idx_new < idx_old:
x[idx_new+1:idx_old+1] = x[idx_new:idx_old]
x[idx_new] = data[i]
else:
x[idx_new] = data[i]
return res
Esto nos lleva a 12 ms, una mejora x60 sobre el enfoque de pandas + SciPy sin procesar:
> %timeit rolling_mean_np_jitted_search(data.values)
12 ms ± 210 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
Podrías intentar usar scipy.stats.trim_mean :
from scipy.stats import trim_mean
df['value'].rolling(5).apply(lambda x: trim_mean(x, 0.2))
[output]
0 NaN
1 NaN
2 NaN
3 NaN
4 10.000000
5 11.000000
6 13.000000
7 13.333333
8 14.000000
9 15.666667
Tenga en cuenta que tuve que usar rolling(5) y proportiontocut=0.2 para su conjunto de datos de juguetes.
Para sus datos reales, debe usar rolling(50) y trim_mean(x, 0.06) para eliminar los 3 valores superior e inferior de la ventana móvil.