| |
| |
| |
| |
| |
|
|
| from torch.utils.data import Dataset, DataLoader |
| import pandas as pd |
| import numpy as np |
| import os |
| from sklearn.preprocessing import StandardScaler |
| from decompositions.decomposition import emd_augment |
|
|
| class Dataset_ETT_hour(Dataset): |
| def __init__(self, root_path, flag='train', size=None, |
| features='S', data_path='ETTh1.csv', |
| target='OT', scale=True, freq='h', n_imf = 500, percentage = 100, params = None): |
| |
| |
| if size == None: |
| self.seq_len = 24 * 4 * 4 |
| self.label_len = 24 * 4 |
| self.pred_len = 24 * 4 |
| else: |
| self.seq_len = size[0] |
| self.label_len = size[1] |
| self.pred_len = size[2] |
| |
| assert flag in ['train', 'test', 'val'] |
| type_map = {'train': 0, 'val': 1, 'test': 2} |
| self.set_type = type_map[flag] |
|
|
| self.features = features |
| self.target = target |
| self.scale = scale |
| self.freq = freq |
| self.n_imf = n_imf |
| self.percentage = percentage |
|
|
| self.root_path = root_path |
| self.data_path = data_path |
| self.params = params |
| self.__read_data__() |
|
|
| def __read_data__(self): |
| self.scaler = StandardScaler() |
| df_raw = pd.read_csv(os.path.join(self.root_path, |
| self.data_path)) |
|
|
|
|
| border1s = [0, 12 * 30 * 24 - self.seq_len, 12 * 30 * 24 + 4 * 30 * 24 - self.seq_len] |
| border2s = [12 * 30 * 24, 12 * 30 * 24 + 4 * 30 * 24, 12 * 30 * 24 + 8 * 30 * 24] |
| border1 = border1s[self.set_type] |
| border2 = border2s[self.set_type] |
|
|
| if self.features == 'M' or self.features == 'MS': |
| cols_data = df_raw.columns[1:] |
| df_data = df_raw[cols_data] |
| elif self.features == 'S': |
| df_data = df_raw[[self.target]] |
|
|
| if self.scale: |
| train_data = df_data[border1s[0]:border2s[0]] |
| train_length = int((self.percentage / 100) * len(train_data)) |
| train_data = train_data[-train_length:] |
| self.scaler.fit(train_data.values) |
| data = self.scaler.transform(df_data.values) |
| else: |
| data = df_data.values |
|
|
| if self.set_type == 0 and self.params.aug_type == 5: |
| self.aug_data = emd_augment(data[border1:border2][-len(train_data):], self.seq_len+self.pred_len, n_IMF = self.n_imf) |
| else: |
| self.aug_data = np.zeros_like(data[border1:border2]) |
| |
| if self.set_type == 0: |
| self.data_x = data[border1:border2][-len(train_data):] |
| self.data_y = data[border1:border2][-len(train_data):] |
| else: |
| self.data_x = data[border1:border2] |
| self.data_y = data[border1:border2] |
|
|
|
|
| def __getitem__(self, index): |
| s_begin = index |
| s_end = s_begin + self.seq_len |
| r_begin = s_end - self.label_len |
| r_end = r_begin + self.label_len + self.pred_len |
|
|
| seq_x = self.data_x[s_begin:s_end] |
| seq_y = self.data_y[r_begin:r_end] |
|
|
| if self.params.aug_type == 5: |
| aug_data = self.aug_data[s_begin] |
| else: |
| aug_data = np.array([]) |
| |
| return seq_x, seq_y, aug_data |
|
|
| def __len__(self): |
| return len(self.data_x) - self.seq_len - self.pred_len + 1 |
|
|
| def inverse_transform(self, data): |
| return self.scaler.inverse_transform(data) |
| |
| class Dataset_ETT_minute(Dataset): |
| def __init__(self, root_path, flag='train', size=None, |
| features='S', data_path='ETTm1.csv', |
| target='OT', scale=True, freq='t', n_imf = 500, percentage = 100, params = None): |
| |
| |
| if size == None: |
| self.seq_len = 24 * 4 * 4 |
| self.label_len = 24 * 4 |
| self.pred_len = 24 * 4 |
| else: |
| self.seq_len = size[0] |
| self.label_len = size[1] |
| self.pred_len = size[2] |
| |
| assert flag in ['train', 'test', 'val'] |
| type_map = {'train': 0, 'val': 1, 'test': 2} |
| self.set_type = type_map[flag] |
|
|
| self.features = features |
| self.target = target |
| self.scale = scale |
| self.freq = freq |
| self.n_imf = n_imf |
|
|
| self.root_path = root_path |
| self.data_path = data_path |
| self.percentage = percentage |
| self.params = params |
| self.__read_data__() |
|
|
| def __read_data__(self): |
| self.scaler = StandardScaler() |
| df_raw = pd.read_csv(os.path.join(self.root_path, |
| self.data_path)) |
|
|
| border1s = [0, 12 * 30 * 24 * 4 - self.seq_len, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4 - self.seq_len] |
| border2s = [12 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 4 * 30 * 24 * 4, 12 * 30 * 24 * 4 + 8 * 30 * 24 * 4] |
| border1 = border1s[self.set_type] |
| border2 = border2s[self.set_type] |
|
|
| if self.features == 'M' or self.features == 'MS': |
| cols_data = df_raw.columns[1:] |
| df_data = df_raw[cols_data] |
| elif self.features == 'S': |
| df_data = df_raw[[self.target]] |
|
|
| if self.scale: |
| train_data = df_data[border1s[0]:border2s[0]] |
| train_length = int((self.percentage / 100) * len(train_data)) |
| train_data = train_data[-train_length:] |
| self.scaler.fit(train_data.values) |
| data = self.scaler.transform(df_data.values) |
| else: |
| data = df_data.values |
|
|
| if self.set_type == 0 and self.params.aug_type == 5: |
| self.aug_data = emd_augment(data[border1:border2][-len(train_data):], self.seq_len+self.pred_len, n_IMF = self.n_imf) |
| else: |
| self.aug_data = np.zeros_like(data[border1:border2]) |
|
|
| if self.set_type == 0: |
| self.data_x = data[border1:border2][-len(train_data):] |
| self.data_y = data[border1:border2][-len(train_data):] |
| else: |
| self.data_x = data[border1:border2] |
| self.data_y = data[border1:border2] |
|
|
| def __getitem__(self, index): |
| s_begin = index |
| s_end = s_begin + self.seq_len |
| r_begin = s_end - self.label_len |
| r_end = r_begin + self.label_len + self.pred_len |
|
|
| seq_x = self.data_x[s_begin:s_end] |
| seq_y = self.data_y[r_begin:r_end] |
|
|
| if self.params.aug_type == 5: |
| aug_data = self.aug_data[s_begin] |
| else: |
| aug_data = np.array([]) |
|
|
| return seq_x, seq_y, aug_data |
| |
| def __len__(self): |
| return len(self.data_x) - self.seq_len - self.pred_len + 1 |
|
|
| def inverse_transform(self, data): |
| return self.scaler.inverse_transform(data) |
|
|
|
|
| class Dataset_Custom(Dataset): |
| def __init__(self, root_path, flag='train', size=None, |
| features='S', data_path='ETTh1.csv', scale = True, |
| target='OT', freq='h', n_imf = 500, percentage = 100, params=None): |
| |
| |
|
|
| self.seq_len = size[0] |
| self.label_len = size[1] |
| self.pred_len = size[2] |
| |
| assert flag in ['train', 'test', 'val'] |
| type_map = {'train': 0, 'val': 1, 'test': 2} |
| self.set_type = type_map[flag] |
|
|
| self.features = features |
| self.target = target |
| self.freq = freq |
| self.scale = scale |
| self.n_imf = n_imf |
| self.percentage = percentage |
| self.root_path = root_path |
| self.data_path = data_path |
| self.params = params |
|
|
| self.__read_data__() |
|
|
| def __read_data__(self): |
|
|
| self.scaler = StandardScaler() |
| df_raw = pd.read_csv(os.path.join(self.root_path, |
| self.data_path)) |
| ''' |
| df_raw.columns: ['date', ...(other features), target feature] |
| ''' |
| cols = list(df_raw.columns) |
| cols.remove(self.target) |
| cols.remove('date') |
| df_raw = df_raw[['date'] + cols + [self.target]] |
|
|
| num_train = int(len(df_raw) * 0.7) |
| num_test = int(len(df_raw) * 0.2) |
| num_vali = len(df_raw) - num_train - num_test |
| border1s = [0, num_train - self.seq_len, len(df_raw) - num_test - self.seq_len] |
| border2s = [num_train, num_train + num_vali, len(df_raw)] |
| border1 = border1s[self.set_type] |
| border2 = border2s[self.set_type] |
|
|
| if self.features == 'M' or self.features == 'MS': |
| cols_data = df_raw.columns[1:] |
| df_data = df_raw[cols_data] |
| elif self.features == 'S': |
| df_data = df_raw[[self.target]] |
|
|
| if self.scale: |
| train_data = df_data[border1s[0]:border2s[0]] |
| train_length = int((self.percentage / 100) * len(train_data)) |
| train_data = train_data[-train_length:] |
| self.scaler.fit(train_data.values) |
| data = self.scaler.transform(df_data.values) |
| else: |
| data = df_data.values |
|
|
| if self.set_type == 0 and self.params.aug_type == 5: |
| self.aug_data = emd_augment(data[border1:border2][-len(train_data):], self.seq_len+self.pred_len, n_IMF = self.n_imf) |
| |
| else: |
| self.aug_data = np.zeros_like(data[border1:border2]) |
|
|
| if self.set_type == 0: |
| self.data_x = data[border1:border2][-len(train_data):] |
| self.data_y = data[border1:border2][-len(train_data):] |
| else: |
| self.data_x = data[border1:border2] |
| self.data_y = data[border1:border2] |
|
|
| def __getitem__(self, index): |
| s_begin = index |
| s_end = s_begin + self.seq_len |
| r_begin = s_end - self.label_len |
| r_end = r_begin + self.label_len + self.pred_len |
|
|
| seq_x = self.data_x[s_begin:s_end] |
| seq_y = self.data_y[r_begin:r_end] |
|
|
| if self.params.aug_type == 5: |
| aug_data = self.aug_data[s_begin] |
| else: |
| aug_data = np.array([]) |
| |
| return seq_x, seq_y, aug_data |
|
|
| def __len__(self): |
| return len(self.data_x) - self.seq_len - self.pred_len + 1 |
|
|
| def inverse_transform(self, data): |
| return self.scaler.inverse_transform(data) |
|
|
|
|
| class Dataset_Pred(Dataset): |
|
|
| def __init__(self, root_path, flag='pred', size=None, |
| features='S', data_path='ETTh1.csv', |
| target='OT', scale=True, inverse=False, freq='15min', cols=None): |
| |
| |
|
|
| self.seq_len = size[0] |
| self.label_len = size[1] |
| self.pred_len = size[2] |
| |
| assert flag in ['pred'] |
|
|
| self.features = features |
| self.target = target |
| self.scale = scale |
| self.inverse = inverse |
| self.freq = freq |
| self.cols = cols |
| self.root_path = root_path |
| self.data_path = data_path |
| self.__read_data__() |
|
|
| def __read_data__(self): |
| self.scaler = StandardScaler() |
| df_raw = pd.read_csv(os.path.join(self.root_path, |
| self.data_path)) |
| ''' |
| df_raw.columns: ['date', ...(other features), target feature] |
| ''' |
| if self.cols: |
| cols = self.cols.copy() |
| cols.remove(self.target) |
| else: |
| cols = list(df_raw.columns) |
| cols.remove(self.target) |
| cols.remove('date') |
| df_raw = df_raw[['date'] + cols + [self.target]] |
| border1 = len(df_raw) - self.seq_len |
| border2 = len(df_raw) |
|
|
| if self.features == 'M' or self.features == 'MS': |
| cols_data = df_raw.columns[1:] |
| df_data = df_raw[cols_data] |
| elif self.features == 'S': |
| df_data = df_raw[[self.target]] |
|
|
| if self.scale: |
| self.scaler.fit(df_data.values) |
| data = self.scaler.transform(df_data.values) |
| else: |
| data = df_data.values |
|
|
|
|
| self.data_x = data[border1:border2] |
| if self.inverse: |
| self.data_y = df_data.values[border1:border2] |
| else: |
| self.data_y = data[border1:border2] |
|
|
| def __getitem__(self, index): |
| s_begin = index |
| s_end = s_begin + self.seq_len |
| r_begin = s_end - self.label_len |
| r_end = r_begin + self.label_len + self.pred_len |
|
|
| seq_x = self.data_x[s_begin:s_end] |
| if self.inverse: |
| seq_y = self.data_x[r_begin:r_begin + self.label_len] |
| else: |
| seq_y = self.data_y[r_begin:r_begin + self.label_len] |
|
|
| aug_data = np.array([]) |
|
|
| return seq_x, seq_y, aug_data |
|
|
|
|
| def __len__(self): |
| return len(self.data_x) - self.seq_len + 1 |
|
|
| def inverse_transform(self, data): |
| return self.scaler.inverse_transform(data) |
|
|
| data_dict = { |
| 'ETTh1': Dataset_ETT_hour, |
| 'ETTh2': Dataset_ETT_hour, |
| 'ETTm1': Dataset_ETT_minute, |
| 'ETTm2': Dataset_ETT_minute, |
| 'custom': Dataset_Custom, |
| } |
|
|
|
|
| def data_provider(args, flag): |
| Data = data_dict[args.data] |
|
|
| if flag == 'test': |
| shuffle_flag = False |
| drop_last = True |
| batch_size = args.batch_size |
| freq = args.freq |
| nIMF = args.nIMF |
|
|
| elif flag == 'pred': |
| shuffle_flag = False |
| drop_last = False |
| batch_size = 1 |
| freq = args.freq |
| Data = Dataset_Pred |
| else: |
| shuffle_flag = True |
| drop_last = True |
| batch_size = args.batch_size |
| freq = args.freq |
| nIMF = args.nIMF |
|
|
|
|
| data_set = Data( |
| root_path=args.root_path, |
| data_path=args.data_path, |
| flag=flag, |
| size=[args.seq_len, args.label_len, args.pred_len], |
| features=args.features, |
| target=args.target, |
| freq=freq, |
| n_imf = nIMF, |
| percentage = args.percentage, |
| params = args |
| ) |
| data_loader = DataLoader( |
| data_set, |
| batch_size=batch_size, |
| shuffle=shuffle_flag, |
| num_workers=args.num_workers, |
| drop_last=drop_last) |
| return data_set, data_loader |
|
|
