WaveLSFromer / data_loader.py

ducheng678

Initial WaveLSFromer project

093b0a5 6 months ago

22.7 kB

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)

	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)