124 lines
5.6 KiB
Python
124 lines
5.6 KiB
Python
"""
|
|
Author: Weisen Pan
|
|
Date: 2023-10-24
|
|
"""
|
|
import pandas as pd
|
|
import numpy as np
|
|
from tqdm import tqdm
|
|
|
|
def find_pos(out_degree_matrix, num_nodes):
|
|
stage = np.zeros(num_nodes)
|
|
signal = True
|
|
while signal:
|
|
temp = stage.copy()
|
|
for m in range(num_nodes):
|
|
for n in range(num_nodes):
|
|
if out_degree_matrix[m, n] != 0:
|
|
stage[n] = max(stage[n], stage[m] + 1)
|
|
if (temp == stage).all():
|
|
signal = False
|
|
return stage
|
|
|
|
def create_position(pos, num_feat):
|
|
pe = np.array([[[posit / (10000.0 ** (i // 2 * 2.0 / num_feat)) for i in range(num_feat)] for posit in posi] for posi in pos])
|
|
pe[:, :, 0::2] = np.sin(pe[:, :, 0::2])
|
|
pe[:, :, 1::2] = np.cos(pe[:, :, 1::2])
|
|
return pe
|
|
|
|
def create_attn_mask(tensor, num_heads, num_nodes):
|
|
mask = np.zeros(((tensor.shape[0] * num_heads), tensor.shape[1], tensor.shape[2]))
|
|
for x in range(0, mask.shape[0], num_heads):
|
|
mask[x:x + num_heads] = tensor[x // num_heads] + np.eye(num_nodes)
|
|
return mask.reshape(tensor.shape[0], num_heads, tensor.shape[1], tensor.shape[2])
|
|
|
|
def prepare_data_exp1_dag(pred_task, pred_tgt, pred_mode):
|
|
direct = './datasets_exp1/%s%s/%s' % (pred_task, pred_tgt, pred_mode)
|
|
df_train = pd.read_csv(direct + '/train.csv')
|
|
df_val = pd.read_csv(direct + '/val.csv')
|
|
df_test = pd.read_csv(direct + '/test.csv')
|
|
dag_direct = './datasets_exp1/%s%s/' % (pred_task, pred_tgt)
|
|
df_dag_train = pd.read_csv(dag_direct + '/train_daginfo.csv')
|
|
df_dag_val = pd.read_csv(dag_direct + '/val_daginfo.csv')
|
|
df_dag_test = pd.read_csv(dag_direct + '/test_daginfo.csv')
|
|
return df_train, df_val, df_test, df_dag_train, df_dag_val, df_dag_test
|
|
|
|
def prepare_data_exp1(pred_task, pred_tgt, pred_mode):
|
|
direct = './datasets_exp1/%s%s/%s' % (pred_task, pred_tgt, pred_mode)
|
|
df_train = pd.read_csv(direct + '/train.csv')
|
|
df_val = pd.read_csv(direct + '/val.csv')
|
|
df_test = pd.read_csv(direct + '/test.csv')
|
|
return df_train, df_val, df_test
|
|
|
|
def prepare_data_exp23_dag(split):
|
|
direct = './datasets_exp2_3/%s' % split
|
|
df_train = pd.read_csv(direct + '/train.csv')
|
|
df_val = pd.read_csv(direct + '/val.csv')
|
|
df_test = pd.read_csv(direct + '/test.csv')
|
|
df_dag_train = pd.read_csv(direct + '/train_daginfo.csv')
|
|
df_dag_val = pd.read_csv(direct + '/val_daginfo.csv')
|
|
df_dag_test = pd.read_csv(direct + '/test_daginfo.csv')
|
|
return df_train, df_val, df_test, df_dag_train, df_dag_val, df_dag_test
|
|
|
|
def prepare_data_exp23(split):
|
|
direct = './datasets_exp2_3/%s' % split
|
|
df_train = pd.read_csv(direct + '/train.csv')
|
|
df_val = pd.read_csv(direct + '/val.csv')
|
|
df_test = pd.read_csv(direct + '/test.csv')
|
|
return df_train, df_val, df_test
|
|
|
|
def preprocess_data_exp1_dag(pred_task, pred_tgt, pred_mode, num_feat=34, num_task=7, num_head=8):
|
|
df_train, df_val, df_test, df_dag_train, df_dag_val, df_dag_test = prepare_data_exp1_dag(pred_task, pred_tgt, pred_mode)
|
|
arr1 = np.array(df_train.iloc[:, :-1])
|
|
arr2 = np.array(df_val.iloc[:, :-1])
|
|
arr3 = np.array(df_test.iloc[:, :-1])
|
|
data = np.vstack((arr1.reshape(arr1.shape[0] * num_task, num_feat),
|
|
arr2.reshape(arr2.shape[0] * num_task, num_feat),
|
|
arr3.reshape(arr3.shape[0] * num_task, num_feat)))
|
|
data = (data - data.min(0)) / (data.max(0) - data.min(0) + 1e-9)
|
|
data1 = data[:arr1.shape[0] * num_task, :]
|
|
data2 = data[arr1.shape[0] * num_task:arr1.shape[0] * num_task + arr2.shape[0] * num_task, :]
|
|
data3 = data[arr1.shape[0] * num_task + arr2.shape[0] * num_task:, :]
|
|
data1 = data1.reshape(-1, num_task, num_feat)
|
|
data2 = data2.reshape(-1, num_task, num_feat)
|
|
data3 = data3.reshape(-1, num_task, num_feat)
|
|
dag1 = df_dag_train.to_numpy().reshape(-1, num_task, num_task * 2 + 1)
|
|
dag2 = df_dag_val.to_numpy().reshape(-1, num_task, num_task * 2 + 1)
|
|
dag3 = df_dag_test.to_numpy().reshape(-1, num_task, num_task * 2 + 1)
|
|
dag1 = dag1[:, :, 1:]
|
|
dag2 = dag2[:, :, 1:]
|
|
dag3 = dag3[:, :, 1:]
|
|
dagout1 = dag1[:, :, :7]
|
|
dagout2 = dag2[:, :, :7]
|
|
dagout3 = dag3[:, :, :7]
|
|
dagin1 = dag1[:, :, 7:]
|
|
dagin2 = dag2[:, :, 7:]
|
|
dagin3 = dag3[:, :, 7:]
|
|
mask1 = dagin1 + dagout1
|
|
mask2 = dagin2 + dagout2
|
|
mask3 = dagin3 + dagout3
|
|
pos1 = np.zeros((arr1.shape[0], 7))
|
|
for x in range(pos1.shape[0]):
|
|
pos1[x] = find_pos(dagout1[x], 7)
|
|
pos2 = np.zeros((arr2.shape[0], 7))
|
|
for x in range(pos2.shape[0]):
|
|
pos2[x] = find_pos(dagout2[x], 7)
|
|
pos3 = np.zeros((arr3.shape[0], 7))
|
|
for x in range(pos3.shape[0]):
|
|
pos3[x] = find_pos(dagout3[x], 7)
|
|
position1 = create_position(pos1, num_feat)
|
|
position2 = create_position(pos2, num_feat)
|
|
position3 = create_position(pos3, num_feat)
|
|
mask_train = create_attn_mask(mask1, num_heads=num_head, num_nodes=num_task)
|
|
mask_val = create_attn_mask(mask2, num_heads=num_head, num_nodes=num_task)
|
|
mask_test = create_attn_mask(mask3, num_heads=num_head, num_nodes=num_task)
|
|
data1 = np.array(data1, dtype=np.float32)
|
|
data2 = np.array(data2, dtype=np.float32)
|
|
data3 = np.array(data3, dtype=np.float32)
|
|
position1 = np.array(position1, dtype=np.float32)
|
|
position2 = np.array(position2, dtype=np.float32)
|
|
position3 = np.array(position3, dtype=np.float32)
|
|
train_data = [(data1[x], df_train.iloc[x, -1], position1[x], mask_train[x]) for x in range(data1.shape[0])]
|
|
val_data = [(data2[x], df_val.iloc[x, -1], position2[x], mask_val[x]) for x in range(data2.shape[0])]
|
|
test_data = [(data3[x], df_test.iloc[x, -1], position3[x], mask_test[x]) for x in range(data3.shape[0])]
|
|
return train_data, val_data, test_data
|