Вот версия numpy.Это дает здоровое (68x на маленьком примере) ускорение.Поскольку он использует линейную корреляцию, этот коэффициент должен стать еще больше, если ваше реальное окно длиннее 3, потому что correlate переключится на более эффективный метод на основе fft.
import numpy as np
import pandas as pd
from scipy import signal
idx = [np.array(['Jan', 'Jan', 'Feb', 'Mar', 'Mar', 'Mar','Apr', 'Apr', 'May', 'Jun', 'Jun', 'Jun','Jul', 'Aug', 'Aug', 'Sep', 'Sep', 'Oct','Oct', 'Oct', 'Nov', 'Dic', 'Dic',]),np.array(['A', 'B', 'B', 'A', 'B', 'C', 'A', 'B', 'B', 'A', 'B', 'C','A', 'B', 'C', 'A', 'B', 'C', 'A', 'B', 'A', 'B', 'C'])]
data = [{'x': 1}, {'x': 5}, {'x': 3}, {'x': 2}, {'x': 7}, {'x': 3},{'x': 1}, {'x': 6}, {'x': 3}, {'x': 5}, {'x': 2}, {'x': 3},{'x': 1}, {'x': 9}, {'x': 3}, {'x': 2}, {'x': 7}, {'x': 3}, {'x': 6}, {'x': 8}, {'x': 2}, {'x': 7}, {'x': 9}]
df = pd.DataFrame(data, index=idx, columns=['x'])
df.index.names=['date','type']
df = df.reset_index()
weights = np.array((0.2,0.3,0.5))
def running_avg():
if 'running' in df.columns:
del df['running']
n = len(weights)
tp, x = df['type'].values, df['x'].values
sidx = np.argsort(tp, kind='stable')
stp = tp[sidx]
bnds = np.where(stp[1:] != stp[:-1])[0] + 1
running = np.empty(sidx.shape)
for bit in np.split(sidx, bnds):
running[bit[:n]] = np.nan
if len(bit) > n:
running[bit[n:]] = signal.correlate(x[bit[:-1]], weights, 'valid', 'auto')
df['running'] = running
def running_OP():
df['rolling']=0
for j in df['type'].unique():
list_1=list(df['x'][df['type']==j])
cumsum = [0]
list_2=list(df['x'][df['type']==j].index)
z=[]
for i, h in enumerate(list_1, 1):
if i>=4:
cumsum.append(0.2*list_1[i-4]+0.3*list_1[i-3]+0.5*list_1[i-2])
else:
cumsum.append('NaN')
cumsum.pop(0)
z.append(cumsum[0])
df['rolling'][list_2]=z
from timeit import repeat
T0 = np.array(repeat(running_OP, repeat=7, number=10))*100
print(f'\nOP: {T0.mean():.3f} ± {T0.std():.3f} ms')
T1 = np.array(repeat(running_avg, repeat=7, number=100))*10000
print(f'pp {T1.mean():.3f} ± {T1.std():.3f} \N{GREEK SMALL LETTER MU}s')
print("\nresults are " + ["different", "equal"][((np.isnan(df['running']) & np.isnan(df['rolling'].astype(float))) | (df['running'] == df['rolling'])).all()])
print(f'speedup roughly {T0.mean()/T1.mean()*1000:.0f}\N{MULTIPLICATION X}')
Пример выполнения:
OP: 62.500 ± 0.473 ms
pp 903.769 ± 11.491 μs
results are equal
speedup roughly 69✕