滚动操作性能缓慢创建新列

Question

您好，我创建的代码在数据框中创建新列时性能非常慢。这个新列是具有不同权重的滚动移动平均值，我从 for-loop 迭代中获得，但我想在 Python 中可能有更好/更快的性能方式来执行此操作。

• 原始数据库与此类似（尽管要大得多）：

  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']
  

它看起来像这样：

           x
date type
Jan  A     1
     B     5
Feb  B     3
Mar  A     2
     B     7
     C     3
Apr  A     1
     B     6
May  B     3
Jun  A     5
     B     2
     C     3
Jul  A     1
Aug  B     9
     C     3
Sep  A     2
     B     7
Oct  C     3
     A     6
     B     8
Nov  A     2
Dic  B     7
     C     9

• 我的目标是改进以下代码以在数据框中创建一个新列（具有不同权重的滚动移动平均值）。我的代码是：

  df=df.reset_index()
  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
  

它看起来像这样：

   date type  x rolling
0   Jan    A  1     NaN
1   Jan    B  5     NaN
2   Feb    B  3     NaN
3   Mar    A  2     NaN
4   Mar    B  7     NaN
5   Mar    C  3     NaN
6   Apr    A  1     NaN
7   Apr    B  6     5.4
8   May    B  3     5.7
9   Jun    A  5     1.3
10  Jun    B  2     4.7
11  Jun    C  3     NaN
12  Jul    A  1     3.2
13  Aug    B  9     3.1
14  Aug    C  3     NaN
15  Sep    A  2     2.2
16  Sep    B  7     5.7
17  Oct    C  3       3
18  Oct    A  6     2.3
19  Oct    B  8     6.6
20  Nov    A  2     3.8
21  Dic    B  7     7.9
22  Dic    C  9       3

** 如果您的代码性能比我的好，那么知道它的速度有多快会很有趣。如果您认为您的代码更好，但您不知道它的速度有多快，请无论如何发布它，因为我会使用更大的数据框来实现。谢谢！

Answer 1

让我们试试这个，看看这是否可以加快你的代码：

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['rolling'] = df.groupby('type')['x'].rolling(4).apply(lambda x: x[-4]*.2 + x[-3]*.3 + x[-2]*.5, raw=True)\
                  .reset_index(level=2, drop=True).swaplevel(0,1)

df

输出：

           x  rolling
date type            
Jan  A     1      NaN
     B     5      NaN
Feb  B     3      NaN
Mar  A     2      NaN
     B     7      NaN
     C     3      NaN
Apr  A     1      NaN
     B     6      5.4
May  B     3      5.7
Jun  A     5      1.3
     B     2      4.7
     C     3      NaN
Jul  A     1      3.2
Aug  B     9      3.1
     C     3      NaN
Sep  A     2      2.2
     B     7      5.7
Oct  C     3      3.0
     A     6      2.3
     B     8      6.6
Nov  A     2      3.8
Dic  B     7      7.9
     C     9      3.0

时间....

您的代码：

每个循环 324 毫秒 ± 1.55 毫秒（7 次运行的平均值 ± 标准偏差，每个循环 1 个）

这段代码：

每个循环 12.6 毫秒 ± 138 微秒（7 次运行的平均值 ± 标准偏差，每次 100 次循环）

Answer 2

这是一个numpy 版本。它提供了健康的加速（在小例子中为 68 倍）。由于它使用线性相关，因此如果您的实际窗口长于 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✕