简化加权移动平均计算:移除.v.to_numpy()直接获取NumPy数组的可行性
简化加权移动平均计算:移除.v.to_numpy()直接获取NumPy数组的可行性
当然可以简化这段代码!你遇到的问题是因为对多列DataFrame调用rolling.apply时,pandas会把函数返回的标量广播到所有列,导致结果是一个所有列内容重复的DataFrame。这里有几个实用的方法可以直接拿到NumPy数组,省去.v.to_numpy()的步骤:
方法1:直接提取结果数组的第一列(不依赖列名)
这种方法最简单,不需要修改原有逻辑,只是把依赖列名的.v.to_numpy()换成更通用的数组索引方式,代码更健壮(就算以后列名变更也不用调整):
accs = (pd.DataFrame({'v': values[start_index:], 'w': weights[start_index:]}) .rolling(window=window_len, min_periods=1, method='table') .apply(weighted_mean, engine='numba', raw=True) .to_numpy()[:, 0][a:b])
.to_numpy()把DataFrame转成二维NumPy数组,[:, 0]取第一列的结果(因为所有列都相同,取任意一列都可以),最后再执行切片[a:b]。
方法2:提前打包数据,让rolling直接返回Series
如果想从根源避免生成多列DataFrame,可以把v和w打包成单个Series的元素,这样rolling.apply的结果会是一个Series而非DataFrame,直接转NumPy数组即可:
# 把v和w组合成二维数组,再转成Series(每个元素是[v, w]的数组) combined_data = pd.Series(np.column_stack([values[start_index:], weights[start_index:]]).tolist()) accs = (combined_data.rolling(window=window_len, min_periods=1, method='table') .apply(lambda window: weighted_mean(np.vstack(window)), engine='numba', raw=True) .to_numpy()[a:b])
这里np.vstack(window)会把窗口内的[v,w]元素重新拼成(window_len, 2)的二维数组,和原来weighted_mean接收的输入格式完全一致,不需要修改你的加权平均函数。
额外提示:关于numba引擎的小细节
因为你用了engine='numba'和raw=True,这已经是效率很高的写法——numba会编译你的weighted_mean函数,避免pandas对象的额外开销。上面两种方法都能保留这个效率优势,不会影响计算速度。
备注:内容来源于stack exchange,提问作者clime
相关产品推荐
相关产品推荐

