如何在Pandas多列DataFrame上计算滚动点积/余弦相似度

核心需求 #

在Pandas DataFrame上实现滚动窗口维度的dot_product（点积）与cosine_similarity（余弦相似度）计算。
目前已知Pandas原生提供df.rolling(window_size).corr()接口可直接计算滚动相关系数，但暂未找到滚动余弦相似度的官方实现方式。

已尝试的错误实现 #

曾尝试调用sklearn的余弦相似度接口配合rolling的table模式实现，代码运行时报错：

from sklearn.metrics.pairwise import cosine_similarity
df.rolling(window=3, method="table").apply(lambda table: cosine_similarity(table.T))

问题复现代码 #

以下代码可完整复现上述报错：

import pandas as pd
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import matplotlib.pyplot as plt
import seaborn as sns
from IPython.display import display
pd.__version__   # 1.4.3

def generate_random_walk(_len):
    lst = [np.random.randn()]
    for i in range(_len): 
        lst.append(lst[i] + np.random.randn())
    return lst

_len = 100
_num_arrays = 5

array_2D = np.array([generate_random_walk(_len) for _ in range(_num_arrays)]).T
df = pd.DataFrame(array_2D)

df.rolling(window=3, method="table").apply(lambda table: cosine_similarity(table.T))

运行环境 #

• Python 3.10
• Pandas 1.4.3

方案优先级要求 #

• 优先使用Pandas原生API实现
• 若无原生对应接口，优先提供规避显式for循环的高效向量化实现：实测Pandas原生corr函数性能比普通for循环实现快100倍以上，普通循环实现性能过差无法满足需求
• 若既无对应Pandas原生函数、也无法实现上述类apply形式的向量化写法，需提供结合numba加速的for循环实现方案，保障计算效率

内容的提问来源于stack exchange，提问作者I. A