Python 当窗口中的所有变量都来自多列时，如何应用滚动函数

我试图计算一个滚动统计，它需要一个窗口中来自两个输入列的所有变量

我唯一的解决方案是使用for循环。有没有更有效的方法，比如使用Pandas的滚动和应用功能

import pandas as pd
from statsmodels.tsa.stattools import coint


def f(x):
    return coint(x['a'], x['b'])[1]


df = pd.DataFrame(data={'a': [1, 2, 3, 4], 'b': [5, 6, 7, 8]})
df2 = df.rolling(2).apply(lambda x: f(x), raw=False)  # KeyError: 'a'

---

我得到KeyError:'a'，因为df一次只传递给f（）一个系列（列）。指定axis=1会将一行和所有列发送到f（），但这两种方法都不能提供所需的观察值集。

您可以尝试滚动、平均和求和：

df['result'] = df.rolling(2).mean().sum(axis=1)

   a  b  result
0  1  5     0.0
1  2  6     7.0
2  3  7     9.0
3  4  8    11.0

编辑

根据OP提出的问题中的新信息添加不同的答案

设置函数

import pandas as pd
from statsmodels.tsa.stattools import coint

def f(x):
    return coint(x['a'], x['b'])

创建数据和数据帧：

a_data = [1,2,3,4]
b_data = [5,6,7,8]

df = pd.DataFrame(data={'a': a_data, 'b': b_data})

   a  b
0  1  5
1  2  6
2  3  7
3  4  8

在研究coint之后，我发现您正在尝试将两个滚动数组传递给f['a']和f['b']。下面将创建数组和数据帧

n=2
arr_a = [df['a'].shift(x).values[::-1][:n] for x in range(len(df['a']))[::-1]]
arr_b = [df['b'].shift(x).values[::-1][:n] for x in range(len(df['b']))[::-1]]

df1 = pd.DataFrame(data={'a': arr_a, 'b': arr_b})

n是滚动窗口的大小

df1
            a           b
0  [1.0, nan]  [5.0, nan]
1  [2.0, 1.0]  [6.0, 5.0]
2  [3.0, 2.0]  [7.0, 6.0]
3      [4, 3]      [8, 7]

然后您可以使用apply.（f）发送数组行

df1.iloc[(n-1):,].apply(f, axis=1)

您的输出如下所示：

1    (-inf, 0.0, [-48.37534, -16.26923, -10.00565])
2    (-inf, 0.0, [-48.37534, -16.26923, -10.00565])
3    (-inf, 0.0, [-48.37534, -16.26923, -10.00565])
dtype: object

当我运行这个程序时，对于完全共线的数据，我确实会得到一个错误，但我怀疑它会随着真实数据而消失

另外，我知道纯矢量化的解决方案可能会更快。我想知道如果这是你想要的，那么这场演出会是什么样子

---

向真正有解决方案的@Zero致敬。

我试着在滚动之前计算总和：

import pandas as pd
import time
df = pd.DataFrame(data={'a': [1, 2, 3, 4], 'b': [5, 6, 7, 8]})
df2 = df.copy()
s = time.time()
df2.loc[:, 'mean1'] = df.sum(axis = 1).rolling(2).mean()
print(time.time() - s)
s = time.time()
df2.loc[:, 'mean2'] = df.rolling(2).mean().sum(axis=1)
print(time.time() - s)
df2
0.003737926483154297
0.005460023880004883
a   b   mean1   mean2
0   1   5   NaN 0.0
1   2   6   7.0 7.0
2   3   7   9.0 9.0
3   4   8   11.0    11.0

它比前面的答案稍微快一点，但效果相同，可能在大型数据集中，差异可能非常显著

您可以修改它以仅选择感兴趣的列：

s = time.time()
print(df[['a', 'b']].sum(axis = 1).rolling(2).mean())
print(time.time() - s)

0     NaN
1     7.0
2     9.0
3    11.0
dtype: float64
0.0033559799194335938

---

谢谢你的回答。它确实解决了最初发布的问题。不过，我已经更新了问题，以便更准确地描述问题。谢谢@run out。有趣的是，将解决方案的性能与矢量化方法（在SO链接中建议）进行比较并没有产生显著差异。