116 lines
4.2 KiB
Python
116 lines
4.2 KiB
Python
import os
|
||
import pandas as pd
|
||
import numpy as np
|
||
from fastdtw import fastdtw
|
||
from tqdm import tqdm
|
||
import torch
|
||
from joblib import Parallel, delayed
|
||
|
||
files = {
|
||
358: ('PeMS03/PEMS03.npz', 'PeMS03/PEMS03.csv'),
|
||
307: ('PeMS04/PEMS04.npz', 'PEMS04.csv'),
|
||
883: ('PeMS07/PEMS07.npz', 'PEMS07.csv'),
|
||
170: ('PeMS08/PEMS08.npz', 'PEMS08.csv')
|
||
}
|
||
|
||
|
||
def compute_dtw_pair(i, j, data_mean):
|
||
return i, j, fastdtw(data_mean[i], data_mean[j], radius=6)[0]
|
||
|
||
|
||
def get_A_hat(args):
|
||
"""Optimized version with GPU support and parallel computing"""
|
||
# 基础配置
|
||
device = torch.device(args['device'])
|
||
data_dir = './data/'
|
||
num_node = args['num_nodes']
|
||
file_npz, file_csv = files[num_node]
|
||
dataset_name = file_npz.split('/')[0]
|
||
os.makedirs(f'{data_dir}{dataset_name}', exist_ok=True)
|
||
|
||
# 数据加载与标准化
|
||
with np.load(f'{data_dir}{file_npz}') as data:
|
||
arr_data = data['data']
|
||
arr_data = (arr_data - arr_data.mean((0, 1))) / arr_data.std((0, 1))
|
||
arr_data = torch.from_numpy(arr_data).float().to(device)
|
||
|
||
# DTW矩阵计算（带缓存）
|
||
dtw_path = f'{data_dir}{dataset_name}/{dataset_name}_dtw_distance.npy'
|
||
if not os.path.exists(dtw_path):
|
||
# 使用GPU加速的均值计算
|
||
daily_data = arr_data[..., 0].unfold(0, 288, 288).mean(dim=0).T.cpu().numpy()
|
||
|
||
# 并行计算DTW
|
||
print("Computing DTW matrix with parallel optimization...")
|
||
results = Parallel(n_jobs=-1)(
|
||
delayed(compute_dtw_pair)(i, j, daily_data)
|
||
for i in tqdm(range(num_node)) for j in range(i, num_node)
|
||
)
|
||
|
||
dtw_matrix = np.full((num_node, num_node), np.inf)
|
||
for i, j, d in results:
|
||
dtw_matrix[i, j] = d
|
||
dtw_matrix[j, i] = d
|
||
np.save(dtw_path, dtw_matrix)
|
||
else:
|
||
dtw_matrix = np.load(dtw_path)
|
||
|
||
# DTW矩阵标准化（GPU加速）
|
||
dtw_tensor = torch.from_numpy(dtw_matrix).to(device)
|
||
dtw_normalized = (dtw_tensor - dtw_tensor.mean()) / dtw_tensor.std()
|
||
semantic_adj = torch.exp(-dtw_normalized ** 2 / args['sigma1'] ** 2)
|
||
semantic_adj = (semantic_adj > args['thres1']).float()
|
||
|
||
# 空间矩阵计算（带缓存）
|
||
spatial_path = f'{data_dir}{dataset_name}/{dataset_name}_spatial_distance.npy'
|
||
if not os.path.exists(spatial_path):
|
||
# 使用Pandas高效读取
|
||
df = pd.read_csv(f'{data_dir}{file_csv}', header=None)
|
||
if num_node == 358: # 特殊处理节点ID映射
|
||
with open(f'{data_dir}{dataset_name}/{dataset_name}.txt') as f:
|
||
node_ids = [int(line.strip()) for line in f]
|
||
id_map = {nid: idx for idx, nid in enumerate(node_ids)}
|
||
df[0] = df[0].map(id_map)
|
||
df[1] = df[1].map(id_map)
|
||
|
||
# 构建稀疏矩阵
|
||
spatial_adj = torch.full((num_node, num_node), float('inf'), device=device)
|
||
for row in df.itertuples():
|
||
i, j, d = int(row[1]), int(row[2]), float(row[3])
|
||
spatial_adj[i, j] = spatial_adj[j, i] = d
|
||
spatial_adj = spatial_adj.cpu().numpy()
|
||
np.save(spatial_path, spatial_adj)
|
||
else:
|
||
spatial_adj = np.load(spatial_path)
|
||
|
||
# 空间矩阵标准化（GPU加速）
|
||
mask = spatial_adj != float('inf')
|
||
valid_values = torch.from_numpy(spatial_adj[mask]).to(device)
|
||
spatial_normalized = (spatial_adj - valid_values.mean().item()) / valid_values.std().item()
|
||
spatial_adj = torch.exp(-torch.tensor(spatial_normalized) ** 2 / args['sigma2'] ** 2)
|
||
spatial_adj = (spatial_adj > args['thres2']).float()
|
||
|
||
# 归一化处理
|
||
def normalize_adj(adj):
|
||
D = adj.sum(1)
|
||
D = torch.clamp(D, min=1e-5)
|
||
D_inv_sqrt = 1.0 / torch.sqrt(D)
|
||
return 0.8 * (torch.eye(adj.size(0), device=device) +
|
||
0.8 * D_inv_sqrt.view(-1, 1) * adj * D_inv_sqrt.view(1, -1))
|
||
|
||
return (normalize_adj(semantic_adj.to(args['device'])), normalize_adj(spatial_adj.to(args['device'])))
|
||
|
||
|
||
# 测试代码
|
||
if __name__ == '__main__':
|
||
config = {
|
||
'sigma1': 0.1,
|
||
'sigma2': 10,
|
||
'thres1': 0.6,
|
||
'thres2': 0.5,
|
||
'device': 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
||
}
|
||
|
||
for nodes in [358, 883, 170]:
|
||
args = {'num_nodes': nodes, **config}
|
||
get_A_hat(args) |