import os import numpy as np import pandas as pd import torch from torch.utils.data import Dataset # from sklearn.preprocessing import StandardScaler from utils.tools import StandardScaler, dotdict from utils.timefeatures import time_features import warnings warnings.filterwarnings("ignore") class Dataset_ETT_hour(Dataset): def __init__( self, root_path, flag="train", size=None, features="S", data_path="ETTh1.csv", target="OT", scale=True, inverse=False, timeenc=0, freq="h", cols=None, ): # size [seq_len, label_len, pred_len] # info 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] # init 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.inverse = inverse self.timeenc = timeenc self.freq = freq 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)) 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]] self.scaler.fit(train_data.values, scale_mean=not self.config.no_scale_mean) data = self.scaler.transform(df_data.values) else: data = df_data.values df_stamp = df_raw[["date"]][border1:border2] df_stamp["date"] = pd.to_datetime(df_stamp.date) self.raw_dates = df_stamp.date.to_numpy(dtype=np.datetime64) data_stamp = time_features(df_stamp, timeenc=self.timeenc, freq=self.freq) self.data_x = data[border1:border2] if self.inverse: self.data_y = df_data.values[border1:border2] else: self.data_y = data[border1:border2] self.data_stamp = data_stamp 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 = np.concatenate( [ self.data_x[r_begin : r_begin + self.label_len], self.data_y[r_begin + self.label_len : r_end], ], 0, ) else: seq_y = self.data_y[r_begin:r_end] seq_x_mark = self.data_stamp[s_begin:s_end] seq_y_mark = self.data_stamp[r_begin:r_end] return seq_x, seq_y, seq_x_mark, seq_y_mark, index def index_to_dates(self, index): # index is of length batch_size s_begin = index s_end = s_begin + self.config.seq_len r_begin = s_end - self.config.label_len r_end = r_begin + self.config.label_len + self.config.pred_len seq_x_raw_dates = self.raw_dates[ np.add.outer(s_begin, np.arange(self.config.seq_len)) ] seq_y_raw_dates = self.raw_dates[ np.add.outer( r_begin, np.arange(self.config.label_len + self.config.pred_len) ) ] return seq_x_raw_dates, seq_y_raw_dates 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, inverse=False, timeenc=0, freq="t", cols=None, ): # size [seq_len, label_len, pred_len] # info 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] # init 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.inverse = inverse self.timeenc = timeenc self.freq = freq 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)) 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]] self.scaler.fit(train_data.values, scale_mean=not self.config.no_scale_mean) data = self.scaler.transform(df_data.values) else: data = df_data.values df_stamp = df_raw[["date"]][border1:border2] df_stamp["date"] = pd.to_datetime(df_stamp.date) self.raw_dates = df_stamp.date.to_numpy(dtype=np.datetime64) data_stamp = time_features(df_stamp, timeenc=self.timeenc, freq=self.freq) self.data_x = data[border1:border2] if self.inverse: self.data_y = df_data.values[border1:border2] else: self.data_y = data[border1:border2] self.data_stamp = data_stamp 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 = np.concatenate( [ self.data_x[r_begin : r_begin + self.label_len], self.data_y[r_begin + self.label_len : r_end], ], 0, ) else: seq_y = self.data_y[r_begin:r_end] seq_x_mark = self.data_stamp[s_begin:s_end] seq_y_mark = self.data_stamp[r_begin:r_end] return seq_x, seq_y, seq_x_mark, seq_y_mark, index def index_to_dates(self, index): # index is of length batch_size s_begin = index s_end = s_begin + self.config.seq_len r_begin = s_end - self.config.label_len r_end = r_begin + self.config.label_len + self.config.pred_len seq_x_raw_dates = self.raw_dates[ np.add.outer(s_begin, np.arange(self.config.seq_len)) ] seq_y_raw_dates = self.raw_dates[ np.add.outer( r_begin, np.arange(self.config.label_len + self.config.pred_len) ) ] return seq_x_raw_dates, seq_y_raw_dates 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, config, flag="train"): # Default values defaults = { "size": None, "features": "S", "target": "OT", "scale": True, "inverse_pred": False, "inverse_output": False, "cols": None, "date_start": None, "date_end": None, "date_test": None, "date_val": None, "t_embed": None, } config = dotdict({**defaults, **config}) assert config.seq_len is not None assert config.label_len is not None assert config.pred_len is not None assert flag in ["train", "test", "val"] assert config.freq is not None assert config.root_path is not None assert config.data_path is not None assert ( (config.date_start is None) or (config.date_end is None) or (config.date_start < config.date_end) ), "date_start isn't before date_end" assert ( (config.date_test is None) or (config.date_end is None) or (config.date_test < config.date_end) ), "date_test isn't before date_end" assert ( (config.date_test is None) or (config.date_start is None) or (config.date_test > config.date_start) ), "date_test isn't after date_start" assert (config.date_val is None) or ( config.date_test is not None ), "date_val is used without date_test" assert ( (config.date_val is None) or (config.date_test is None) or (config.date_val < config.date_test) ), "date_val isn't before date_test" assert ( (config.date_val is None) or (config.date_end is None) or (config.date_val < config.date_end) ), "date_val isn't before date_end" assert ( (config.date_val is None) or (config.date_start is None) or (config.date_val > config.date_start) ), "date_val isn't after date_start" assert (config.label_len == 0) or ( config.inverse_output == config.inverse_pred ), "If label length is non-zero then inverse_pred and inverse_output should be the same" self.config = config self.flag = flag # self.timeenc = 0 if config.t_embed != "timeF" else 1 if config.t_embed == "timeF": self.timeenc = 1 elif config.t_embed == "time2vec_add" or config.t_embed == "time2vec_app": self.timeenc = 2 else: self.timeenc = 0 type_map = {"train": 0, "val": 1, "test": 2} self.set_type = type_map[flag] self.__read_data__() def __read_data__(self): self.scaler = StandardScaler() df_raw = pd.read_csv(os.path.join(self.config.root_path, self.config.data_path)) df_raw = df_raw.astype( {c: np.float32 for c in df_raw.select_dtypes(include="float64").columns} ) df_raw["date"] = pd.to_datetime(df_raw["date"]) if np.isinf(df_raw[df_raw.columns[1:]].to_numpy()).any(): raise Exception("There are inf's in the dataset") if np.isnan(df_raw[df_raw.columns[1:]].to_numpy()).any(): raise Exception("There are nan's in the dataset") """ df_raw.columns: ['date', ...(other features), target feature] """ # Filter to datapoints in [date_start, date_end] if self.config.date_start is not None: df_raw = df_raw.loc[(df_raw["date"] >= self.config.date_start)] if self.config.date_end is not None: df_raw = df_raw.loc[(df_raw["date"] <= self.config.date_end)] if self.config.cols: cols = self.config.cols.copy() assert self.config.target in cols, "Target not in cols" cols.remove(self.config.target) else: cols = list(df_raw.columns) assert self.config.target in cols, "Target not in data" cols.remove(self.config.target) assert "date" in cols, "`date` not in data" cols.remove("date") df_raw = df_raw[["date"] + cols + [self.config.target]] # Define lengths of train, val, and test if self.config.date_test is not None and self.config.date_val is not None: # num_test and num_val are specified num_test = len(df_raw.loc[df_raw["date"] >= self.config.date_test]) num_vali = len( df_raw.loc[ (df_raw["date"] >= self.config.date_val) & (df_raw["date"] < self.config.date_test) ] ) num_train = len(df_raw) - num_vali - num_test elif self.config.date_test is not None: # num_val is half of num_test which is specified num_test = len(df_raw.loc[(df_raw["date"] >= self.config.date_test)]) num_vali = num_test // 2 num_train = len(df_raw) - num_vali - num_test else: # Default split print("Warning: using default dataset split") num_train = int(len(df_raw) * 0.9) num_test = int(len(df_raw) * 0.05) num_vali = len(df_raw) - num_train - num_test if num_test == 0: raise Exception("Dataset loading issue: num_test==0, check date settings") elif num_vali == 0: raise Exception("Dataset loading issue: num_vali==0, check date settings") elif num_train == 0: raise Exception("Dataset loading issue: num_train==0, check date settings") border1s = [ 0, num_train - self.config.seq_len, len(df_raw) - num_test - self.config.seq_len, ] border2s = [num_train, num_train + num_vali, len(df_raw)] border1 = border1s[self.set_type] border2 = border2s[self.set_type] if self.config.features == "M" or self.config.features == "MS": cols_data = df_raw.columns[1:] df_data = df_raw[cols_data] elif self.config.features == "S": df_data = df_raw[[self.config.target]] if self.config.scale: train_data = df_data[border1s[0] : border2s[0]] self.scaler.fit(train_data.values, scale_mean=not self.config.no_scale_mean) data = torch.from_numpy(self.scaler.transform(df_data.values)) else: data = torch.from_numpy(df_data.values) df_stamp = df_raw[["date"]][border1:border2] df_stamp["date"] = pd.to_datetime(df_stamp.date) self.raw_dates = df_stamp.date.to_numpy(dtype=np.datetime64) self.data_stamp = np.float32( time_features(df_stamp, timeenc=self.timeenc, freq=self.config.freq) ) self.data_x = data[border1:border2] if self.config.inverse_pred: self.data_y = torch.from_numpy(df_data.values[border1:border2]) else: self.data_y = data[border1:border2] def __getitem__(self, index): s_begin = index s_end = s_begin + self.config.seq_len r_begin = s_end - self.config.label_len r_end = r_begin + self.config.label_len + self.config.pred_len seq_x = self.data_x[s_begin:s_end] if self.config.inverse_pred: # this is where inverse_pred != inverse output gets wonky if label_len != 0 # its because the label doesn't get inversed seq_y = np.concatenate( [ self.data_x[ r_begin : r_begin + self.config.label_len ], # Use non-scaled data_x self.data_y[r_begin + self.config.label_len : r_end], ], axis=0, ) else: seq_y = self.data_y[r_begin:r_end] seq_x_mark = self.data_stamp[s_begin:s_end] seq_y_mark = self.data_stamp[r_begin:r_end] return seq_x, seq_y, seq_x_mark, seq_y_mark, index def index_to_dates(self, index): # index is of length batch_size s_begin = index s_end = s_begin + self.config.seq_len r_begin = s_end - self.config.label_len r_end = r_begin + self.config.label_len + self.config.pred_len seq_x_raw_dates = self.raw_dates[ np.add.outer(s_begin, np.arange(self.config.seq_len)) ] seq_y_raw_dates = self.raw_dates[ np.add.outer( r_begin, np.arange(self.config.label_len + self.config.pred_len) ) ] # seq_x_raw_dates = self.raw_dates[np.r_[s_begin,s_end-1].reshape(-1, index.shape[0]).T]# self.raw_dates.iloc[np.r_[s_begin,s_end]] # seq_y_raw_dates = self.raw_dates[np.r_[r_begin,r_end-1].reshape(-1, index.shape[0]).T]# self.raw_dates.iloc[np.r_[r_begin,r_end]] return seq_x_raw_dates, seq_y_raw_dates def __len__(self): return len(self.data_x) - self.config.seq_len - self.config.pred_len + 1 def inverse_transform(self, data): return self.scaler.inverse_transform(data) class Dataset_Pred(Dataset): def __init__(self, config, flag="pred"): # Default values defaults = { "size": None, "features": "S", "target": "OT", "scale": True, "inverse": False, "cols": None, "date_start": None, "date_end": None, "t_embed": None, } config = dotdict({**defaults, **config}) assert config.seq_len is not None assert config.label_len is not None assert config.pred_len is not None assert flag in ["pred"] assert config.freq is not None assert config.root_path is not None assert config.data_path is not None assert ( (config.date_start is None) or (config.date_end is None) or (config.date_start < config.date_end) ), "date_start isn't before date_end" self.config = config self.flag = flag # self.timeenc = 0 if config.t_embed != "timeF" else 1 if config.t_embed == "timeF": self.timeenc = 1 elif config.t_embed == "time2vec_add" or config.t_embed == "time2vec_app": self.timeenc = 2 else: self.timeenc = 0 self.__read_data__() def __read_data__(self): self.scaler = StandardScaler() df_raw = pd.read_csv(os.path.join(self.config.root_path, self.config.data_path)) """ df_raw.columns: ['date', ...(other features), target feature] """ # Filter to datapoints in [date_start, date_end] if self.config.date_start is not None: df_raw = df_raw.loc[(df_raw["date"] >= self.config.date_start)] if self.config.date_end is not None: df_raw = df_raw.loc[(df_raw["date"] <= self.config.date_end)] if self.config.cols: cols = self.config.cols.copy() cols.remove(self.config.target) else: cols = list(df_raw.columns) cols.remove(self.config.target) cols.remove("date") df_raw = df_raw[["date"] + cols + [self.config.target]] border1 = len(df_raw) - self.config.seq_len border2 = len(df_raw) if self.config.features == "M" or self.config.features == "MS": cols_data = df_raw.columns[1:] df_data = df_raw[cols_data] elif self.config.features == "S": df_data = df_raw[[self.config.target]] if self.config.scale: self.scaler.fit(df_data.values, scale_mean=not self.config.no_scale_mean) data = self.scaler.transform(df_data.values) else: data = df_data.values tmp_stamp = df_raw[["date"]][border1:border2] tmp_stamp["date"] = pd.to_datetime(tmp_stamp.date) pred_dates = pd.date_range( tmp_stamp.date.values[-1], periods=self.config.pred_len + 1, freq=self.config.freq, ) df_stamp = pd.DataFrame(columns=["date"]) df_stamp.date = pd.to_datetime( list(tmp_stamp.date.values) + list(pred_dates[1:]), utc=True ) self.raw_dates = df_stamp.date.to_numpy(dtype=np.datetime64) # TODO: What is the deal with .freq[-1:] self.data_stamp = np.float32( time_features(df_stamp, timeenc=self.timeenc, freq=self.config.freq[-1:]) ) self.data_x = data[border1:border2] if self.config.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.config.seq_len r_begin = s_end - self.config.label_len r_end = r_begin + self.config.label_len + self.config.pred_len seq_x = self.data_x[s_begin:s_end] if self.config.inverse: seq_y = self.data_x[r_begin : r_begin + self.config.label_len] else: seq_y = self.data_y[r_begin : r_begin + self.config.label_len] seq_x_mark = self.data_stamp[s_begin:s_end] seq_y_mark = self.data_stamp[r_begin:r_end] return seq_x, seq_y, seq_x_mark, seq_y_mark, index def index_to_dates(self, index): # index is of length batch_size s_begin = index s_end = s_begin + self.config.seq_len r_begin = s_end - self.config.label_len r_end = r_begin + self.config.label_len + self.config.pred_len seq_x_raw_dates = self.raw_dates[ np.add.outer(s_begin, np.arange(self.config.seq_len)) ] seq_y_raw_dates = self.raw_dates[ np.add.outer( r_begin, np.arange(self.config.label_len + self.config.pred_len) ) ] return seq_x_raw_dates, seq_y_raw_dates def __len__(self): return len(self.data_x) - self.config.seq_len + 1 def inverse_transform(self, data): return self.scaler.inverse_transform(data)