write in front
In recent years, graph neural networks have played an important role in the field of time series forecasting. Among them, Graph Attention Network (GAT) is a graph neural network based on the attention mechanism, which can capture the complex relationships between nodes in graph-structured data, thus achieving outstanding performance in many fields. In this article, we use Pytorch and PyG (PyTorch Geometric is a powerful and flexible graph neural network library) framework to implement a simple example of applying GAT to stock price prediction .
1
Preface
As the complexity of financial markets continues to increase, forecasting stock prices has become a huge challenge. Traditional time series analysis methods, while effective in certain scenarios, often fail to capture the complex interactions and hidden patterns in the market . In order to solve this problem, this paper adopts a new perspective: converting stock price time series into a graph structure, and modeling and analyzing it from a graph perspective through the Graph Attention Network (GAT). The graph attention network is a powerful graph neural network structure that can flexibly capture the relationships between nodes in the graph by introducing an attention mechanism. In the scenario of stock price prediction, the interactions and dependencies between stocks can be naturally modeled as a graph structure , where stocks serve as nodes and the interactions between them serve as edges. Using PyTorch Geometric (PyG), an advanced graph neural network library, this article shows how to build, train, and evaluate a GAT model to predict stock prices.
The core of the project is to convert stock price time series data into a graph structure, and then use GAT's powerful representation capabilities for analysis and prediction. This project is intended as a simple demonstration to show how to use graph attention network to analyze stock price time series. It is only for learning and research purposes and should not be used for actual investment decisions .
2
Environment configuration
Local environment:
Python 3.7
IDE:PycharmLibrary version:
numpy 1.18.1
pandas 1.0.3
torch-geometric 2.0.2
matplotlib 3.2.1
torch 1.10.1
tushare 1.2.603
Code
overall design
Organized through four main documents, this project shows how to use Graph Attention Networks (GAT) for stock price prediction. The entire process is clearly implemented, from acquiring and processing stock data to building and training deep learning models to the final prediction stage. Each file focuses on a specific task, and together they build a complete stock price prediction solution:
1. data_hander.pyData processing module
This module is responsible for getting the stock's closing price from Tushare and saving it to a file. The main functions include: Get the closing price data of the specified stock code and date range. Save and load closing price data. Draw a stock price curve. Construct adjacency matrices between stocks for use in graph models.
2. model.py- Model module
In this module, a deep learning model based on graph attention network (GAT) is defined. The main components include the construction GATPredictorclass, which is used to build the GAT model and set different hidden layer sizes and the number of attention heads. The forward propagation process of the model is defined.
3. train.py- Training module
This file includes functions related to data preprocessing and model training: data normalization and splitting into training and test sets. Use the sliding window method to process time series data. Define the main functions for training and prediction.
4. main.py- Main program module
Serving as the entry point to the project, this file organizes and calls the functions in the three files mentioned above: Defines the ticker and date range to be forecasted. Call related functions for data processing, model building, training and prediction. Defines the main process and execution logic of the entire project.
Data processing module
The data processing module is responsible for obtaining stock price data from external sources and performing appropriate preprocessing for use by the model. Among them, you need to replace tushare API token with Your Token in the code.
def fetch_close_prices(stocks, start_date, end_date, file_name='close_prices.csv'):
if os.path.exists(file_name):
close_prices = pd.read_csv(file_name).values
else:
ts.set_token('Your Token')
pro = ts.pro_api()
close_prices_list = []
for stock in stocks:
df = pro.daily(ts_code=stock, start_date=start_date, end_date=end_date)
close_prices_list.append(df['close'].values)
close_prices = np.column_stack(close_prices_list)
pd.DataFrame(close_prices).to_csv(file_name, index=False)
return close_pricesIn order to use the Graph Attention Network (GAT), the stock data needs to be converted into a graph structure. Here, the Pearson coefficient of the closing price of each stock is used as the basis for constructing an adjacency matrix to represent the correlation between stocks, and a threshold is set to determine whether there are connected edges.
def build_adjacency_matrix(close_prices, threshold=0.5):
N = close_prices.shape[1]
adj_matrix = np.zeros((N, N))
for i in range(N):
for j in range(N):
correlation = np.corrcoef(close_prices[:, i], close_prices[:, j])[0, 1]
adj_matrix[i, j] = 1 if abs(correlation) > threshold else 0
return adj_matrixmodel module
This module is responsible for defining and implementing a simple stock price prediction model based on graph attention network. By constructing a structure including graph attention layer, hidden layer and output layer, the forward pass of the model is realized. Among them, since the GAT layer of PyG can only accept two-dimensional data, in order to improve the parallel computing speed, a batch of data is synthesized into a large picture for calculation, and finally the predicted value is output through the linear layer after conversion through view.
class GATPredictor(nn.Module):
def __init__(self, node_features, node_nums, hidden_size=32, num_heads=1):
super(GATPredictor, self).__init__()
self.node_nums = node_nums
self.gat1 = GATConv(node_features, hidden_size, heads=num_heads)
self.gat2 = GATConv(hidden_size * num_heads, hidden_size, concat=False)
self.out = nn.Linear(hidden_size, 1)
def forward(self, data):
x, edge_index, batch = data.x, data.edge_index, data.batch
batch_size = torch.max(batch).item() + 1
x = torch.relu(self.gat1(x, edge_index))
x = torch.relu(self.gat2(x, edge_index))
x = x.view(batch_size, self.node_nums, -1)
x = self.out(x)
return x.squeeze(2)training module
The training module is mainly responsible for processing the specified input form of data to the model and the logic of model training. It implements some functions for normalizing data and dividing sliding windows, as well as functions for model training and testing.
def normalize_and_split(closing_prices, test_size=0.2):
scaler = MinMaxScaler()
normalized_data = scaler.fit_transform(closing_prices)
train_data, test_data = train_test_split(normalized_data, test_size=test_size, shuffle=False)
return train_data, test_data, scaler
def sliding_window(data, window_size):
windows = []
for i in range(len(data) - window_size):
x = data[i:i + window_size]
y = data[i + window_size]
windows.append((x, y))
return windowWhen constructing the edges between nodes, each stock is used as a node, and the data of its window is used as the node feature, and the same graph relationship is constructed for each sliding window, and finally converted into the specified input data form of PyG.
def create_graph_data(windows, adj_matrix):
graph_data = []
edge_index = torch.tensor(np.where(adj_matrix != 0), dtype=torch.long)
for window in windows:
x, y = window
x_tensor = torch.tensor(x.T, dtype=torch.float)
y_tensor = torch.tensor(y, dtype=torch.float)
data = Data(x=x_tensor, y=y_tensor, edge_index=edge_index)
graph_data.append(data)
return graph_data
def train(model, train_loader, optimizer, criterion, epochs):
model.train()
for epoch in range(epochs):
for batch in train_loader:
optimizer.zero_grad()
out = model(batch)
y = batch.y.view(out.size())
loss = criterion(out, y)
loss.backward()
optimizer.step()
print(f'Epoch: {epoch+1}, Loss: {loss.item()}'Main program module
This module obtains historical closing price data based on a given stock code and date range. Here we simply retrieve several stocks from the Shanghai Stock Exchange in 2022. It is then converted into graph data format through sliding window method and adjacency matrix. Next, initialize and train a prediction model based on the graph attention network, and finally use the model to make predictions and visually compare the prediction results with the real values.
# 股票代码和日期范围
stocks = ['000001.SZ', '000002.SZ', '000006.SZ', '000005.SZ', '000008.SZ', '000009.SZ']
start_date = "2022-01-01"
end_date = "2023-01-01"
window_size = 10
# 获取收盘价并创建邻接矩阵
close_prices = fetch_close_prices(stocks, start_date, end_date)
adj_matrix = build_adjacency_matrix(close_prices)
# 数据归一化和划分
train_data, test_data, scaler = normalize_and_split(close_prices)
# 滑动窗口数据准备
train_windows = sliding_window(train_data, window_size)
test_windows = sliding_window(test_data, window_size)
# 转换为图数据
train_dataset = create_graph_data(train_windows, adj_matrix)
test_dataset = create_graph_data(test_windows, adj_matrix)
print(adj_matrix)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
# 创建模型
model = GATPredictor(node_features=window_size, node_nums=len(stocks))
# 优化器和损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
# 训练
train(model, train_loader, optimizer, criterion, epochs=50)
# 预测
predictions, groundtruth = predict(model, test_loader, scaler)
# 可视化预测结果
plot_results(predictions, groundtruth, stocks)running result
Next, the model is trained and tested. After multiple epoch iterations, the mse gradually converges.
Epoch: 1, Loss: 0.3193877637386322
Epoch: 2, Loss: 0.21741190552711487
Epoch: 3, Loss: 0.16308942437171936
Epoch: 4, Loss: 0.09306434541940689
Epoch: 5, Loss: 0.052657563239336014
Epoch: 6, Loss: 0.02670634351670742
Epoch: 7, Loss: 0.019804660230875015
Epoch: 8, Loss: 0.016364283859729767
Epoch: 9, Loss: 0.017054198309779167
Epoch: 10, Loss: 0.01848340593278408
...
Epoch: 44, Loss: 0.009028978645801544
Epoch: 45, Loss: 0.013259928673505783
Epoch: 46, Loss: 0.01599966734647751
Epoch: 47, Loss: 0.011484017595648766
Epoch: 48, Loss: 0.010215471498668194
Epoch: 49, Loss: 0.01579277403652668
Epoch: 50, Loss: 0.011635522358119488After visualizing the prediction results of multiple stocks, it can be seen that some stocks have better trend fitting effects. Since these six stocks are simply selected, their correlation may not be very strong, so you can also select multiple stocks in the same sector for experiments. In addition, the structure of the model is also very simple and does not introduce more stock price characteristics, so there are many points that can be further improved.
4
Summarize
In the field of modern finance, with the increasing maturity of big data and machine learning technology, traditional stock price prediction methods are facing unprecedented challenges and opportunities. Traditional stock price predictions often rely on the historical data of a single stock and ignore the interaction and impact between stocks. This article explores how to use graph attention network (GAT) to mine potential relationships between stocks, providing a new perspective for stock price prediction. The adjacency matrix is combined with the sliding window method to convert time series stock price data into a graph data format to capture the complex interactions between stocks. The experiment in this article is just a simple demo, and there is still a lot of room for improvement. Interested readers can do further research.
The content of this article is only for technical discussion and learning and does not constitute any investment advice.
To obtain the complete code and data as well as the complete source code and data of other historical articles, you can join the "Artificial Intelligence Quantification Laboratory" Knowledge Planet.
Recommended reading in previous issues
WWW 2023 | Papers related to quantitative trading (with paper links)
KDD 2023 | Papers related to quantitative trading (with paper links)
AAAI 2022 | Papers related to quantitative trading (with paper links)
IJCAI 2022 | Papers related to quantitative trading (with paper links)
WWW 2022 | Papers related to quantitative trading (with paper links)
KDD 2022 | Papers related to quantitative trading (with paper links)
Interpretation: Application of ChatGPT in stock market prediction
Interpretation: What kind of metrics should be used for machine learning prediction revenue models?
[Python Quantification] Deep learning model quantification backtesting framework based on backtrader
[Python Quantification] Use Transformer model for stock price prediction
[Python Quantification] Build a CNN-LSTM model for stock price prediction
[Python Quantification] Use Python to build a stock public opinion analysis system
[Python Quantification] Use Informer for stock price prediction
[python quantification] Using DeepAR for multi-step probability prediction of stock prices
"Artificial Intelligence Quantitative Laboratory" Knowledge Planet
By joining the Knowledge Planet of the Artificial Intelligence Quantitative Laboratory, you can get: (1) Regular push of the latest research results related to the quantitative application of artificial intelligence, including high-level journal papers and high-quality financial engineering research reports from securities companies, so that you can understand the latest cutting-edge knowledge anytime and anywhere; (2) The complete source code of the Python project in the official account’s historical articles; (3) High-quality Python, machine learning, and quantitative trading related e-books PDF; (4) High-quality quantitative trading information and project code sharing; (5) Communicate and make friends with star friends Like-minded friends. (6) Ask the bloggers questions and answer questions.
