高效计算股票相关性：避免重复计算的Python技巧

2024-12-31 23:37:56

高效计算股票相关性：避免重复计算

问题背景

针对美国前500只股票，需要计算它们之间的相关性。直接使用嵌套循环进行两两配对会导致大量重复计算。例如，计算完股票 A 与股票 B 的相关性之后，程序又会重复计算股票 B 与股票 A 的相关性，这种冗余操作消耗计算资源并且降低了执行效率。此类问题在处理大型数据集时尤为突出。如何避免这些重复计算，成为提升效率的关键。

问题分析

重复计算的根本原因是无序的遍历所有可能的配对。在相关性计算中，股票 A 和股票 B 与股票 B 和股票 A 的相关性是相同的。因此，我们只需要计算每个独特的股票配对即可。一个简单但有效的策略是使用一种有规律的方法进行配对，比如只取上三角矩阵或下三角矩阵的方式来计算，从而消除对称的配对。

解决方案一：优化遍历方法

此方案改变循环的方式，仅遍历一半的组合，即通过使用一个辅助的起始点来跳过已经计算过的股票组合，这样确保每个股票组合只被计算一次，有效防止重复计算，从而减少了不必要的开销。

操作步骤:

修改外层循环，使用 enumerate() 来获得当前股票在 tickers 列表中的索引值 i 。
在内层循环中，让内循环的起始索引设置为 i + 1，从而确保它始终在当前外循环索引之后。这样可有效跳过那些先前已经计算过相关性的股票对。

代码示例:

import numpy as np
import pandas as pd

index = 0
ticker_col = np.array([])
compare_col = np.array([])
correlation_col = np.array([])

tickers = [ "AAPL.csv", "MSFT.csv","GOOG.csv",  "AMZN.csv"] #使用一些示例股票进行测试


for i, ticker in enumerate(tickers):
    ticker = ticker[:-1]
    for j in range(i + 1, len(tickers)): # 改变内循环
      
        compare = tickers[j][:-1] #取得正确的比较ticker值
        index += 1


        compare_list = pd.read_csv('C:/stockpredictions/stocks_dfs/' + compare + '.csv')['Adj Close']
        ticker_list = pd.read_csv('C:/stockpredictions/stocks_dfs/' + ticker + '.csv')['Adj Close']
        
        df = pd.DataFrame({'Data1':ticker_list, 'Data2':compare_list}) #准备计算相关性的dataframe
        
        correlation = df['Data1'].corr(df['Data2']) #pearsons r correlation formula

        ticker_col = np.append(ticker_col, ticker) #添加进各个列表方便保存
        compare_col = np.append(compare_col, compare) #添加进各个列表方便保存
        correlation_col = np.append(correlation_col, correlation) #添加进各个列表方便保存

df = pd.DataFrame({'ticker': ticker_col, 'compare': compare_col, 'correlation': correlation_col})
print(df)

原理:

通过限定内层循环的起始索引，保证了股票配对仅进行一次。

安全建议:

注意检查输入数据的合法性，例如股票数据文件是否确实存在。避免出现数据读取的错误，以及后续相关性计算时的错误。

解决方案二：使用 `itertools` 库

itertools 模块提供了高效的迭代器工具。其中， combinations 函数可以方便地生成元素的所有组合，且避免生成重复项。利用它能简洁有效地获得所有独特的股票配对。

操作步骤:

导入 itertools 模块。
使用 itertools.combinations 函数，传入股票代码列表和配对数目（本例为 2），即可获得所有不重复的股票配对。
使用迭代方式来获得股票配对，计算相关性并存储结果。

代码示例:

import numpy as np
import pandas as pd
import itertools

index = 0
ticker_col = np.array([])
compare_col = np.array([])
correlation_col = np.array([])

tickers = [ "AAPL.csv", "MSFT.csv","GOOG.csv",  "AMZN.csv"] # 使用示例股票

for ticker, compare in itertools.combinations(tickers, 2):

    ticker = ticker[:-1]
    compare = compare[:-1]
    
    index += 1

    compare_list = pd.read_csv('C:/stockpredictions/stocks_dfs/' + compare + '.csv')['Adj Close']
    ticker_list = pd.read_csv('C:/stockpredictions/stocks_dfs/' + ticker + '.csv')['Adj Close']

    df = pd.DataFrame({'Data1':ticker_list, 'Data2':compare_list})
    correlation = df['Data1'].corr(df['Data2'])


    ticker_col = np.append(ticker_col, ticker) #添加进各个列表方便保存
    compare_col = np.append(compare_col, compare) #添加进各个列表方便保存
    correlation_col = np.append(correlation_col, correlation) #添加进各个列表方便保存

df = pd.DataFrame({'ticker': ticker_col, 'compare': compare_col, 'correlation': correlation_col})

print(df)