Одно замечание, которое может пригодиться, заключается в том, что вам не нужно сортировать все значения на каждом шаге. Скорее, если вы гарантируете, что окно всегда сортируется, все, что вам нужно сделать, это вставить новое значение в соответствующее место и удалить старое, где оно было, обе из которых являются операциями, которые можно выполнить в O (log_2 (window_size)) используя bisect. На практике это будет выглядеть примерно так:
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
Теперь, дополнительное преимущество в этом случае оказывается меньше, чем то, что дает векторизация, на которую опирается scipy.stats.trim_mean, и, в частности, это все равно будет медленнее, чем решение @ ChrisA, но это полезно отправная точка для дальнейшей оптимизации производительности.
> 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)
Примечательно, что джиттер Нумбы, который часто бывает полезен в подобных ситуациях, также не дает никакой выгоды:
> 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)
Следующий, казалось бы, далеко от оптимального, подход превосходит оба других подхода, рассмотренных выше:
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
Сроки:
> %timeit rolling_mean_np(data.values)
564 ms ± 4.44 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
Более того, на этот раз компиляция JIT делает помощь:
> 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)
Пока мы это делаем, давайте просто быстро проверим, что это действительно делает то, что мы ожидаем:
> np.all(rolling_mean_np_jit(data.values)[49:] == data.rolling(50).apply(lambda w: trim_mean(w, 0.06)).values[49:])
True
На самом деле, помогая сортировщику чуть-чуть, мы можем выжать еще один фактор 2, сократив общее время до 57 мс:
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
Теперь, "сортировка", которая происходит в этом примере, конечно, сводится к тому, чтобы поместить новый элемент в нужное место, в то же время перемещая все между ними на единицу. Выполнение этого вручную сделает чистый код Python более медленным, но версия с джитами приобретает еще один коэффициент 2, что занимает у нас менее 30 мс:
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
На данный момент большую часть времени проводят в np.searchsorted, поэтому давайте сделаем сам поиск JIT-дружественным. Приняв исходный код для bisect, допустим
@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
Это отнимает у нас до 12 мс, что в x60 больше, чем у необработанных панд + SciPy.
> %timeit rolling_mean_np_jitted_search(data.values)
12 ms ± 210 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)