| import os
|
| import datetime
|
| import numpy as np
|
| import pandas as pd
|
|
|
| from torch.utils.data import Dataset
|
| from sklearn.preprocessing import StandardScaler
|
|
|
|
|
|
|
| def _split_with_overlap(data: np.ndarray, train_ratio: float, val_ratio: float, seq_len: int):
|
| """
|
| Time split with overlap for val/test to allow past context:
|
| train: [0 : train_end)
|
| val : [train_end - seq_len : val_end)
|
| test : [val_end - seq_len : T)
|
| """
|
| T = len(data)
|
| train_end = int(T * train_ratio)
|
| val_end = int(T * (train_ratio + val_ratio))
|
|
|
| train_end = max(0, min(train_end, T))
|
| val_end = max(train_end, min(val_end, T))
|
|
|
| val_start = max(0, train_end - seq_len)
|
| test_start = max(0, val_end - seq_len)
|
|
|
| train_data = data[:train_end]
|
| val_data = data[val_start:val_end]
|
| test_data = data[test_start:]
|
|
|
| return train_data, val_data, test_data
|
|
|
|
|
| def _fit_transform_splits(train_data, val_data, test_data, type_flag: str, scaler=None):
|
| if type_flag == "1":
|
| if scaler is None:
|
| scaler = StandardScaler()
|
| scaler.fit(train_data)
|
| train_data = scaler.transform(train_data)
|
| val_data = scaler.transform(val_data)
|
| test_data = scaler.transform(test_data)
|
| return train_data, val_data, test_data, scaler
|
| else:
|
| return train_data, val_data, test_data, None
|
|
|
|
|
| def _to_float32(x: np.ndarray) -> np.ndarray:
|
| return np.asarray(x, dtype=np.float32)
|
|
|
|
|
| def _clean_numeric_csv(df: pd.DataFrame) -> np.ndarray:
|
| """
|
| Keep only numeric columns, and drop common junk index columns.
|
| """
|
| drop_cols = [c for c in df.columns if str(c).lower().startswith("unnamed")]
|
| if drop_cols:
|
| df = df.drop(columns=drop_cols, errors="ignore")
|
|
|
| num_df = df.select_dtypes(include=[np.number])
|
|
|
| if num_df.shape[1] == 0:
|
| raise ValueError("No numeric columns found in CSV after cleaning. Check your file format.")
|
|
|
| num_df = num_df.dropna(axis=0, how="any")
|
|
|
| return num_df.values.astype(np.float32)
|
|
|
|
|
|
|
| class _BaseTimeSeriesDataset(Dataset):
|
|
|
| def __init__(self, flag, seq_len, pre_len):
|
| assert flag in ["train", "val", "test"]
|
| self.flag = flag
|
| self.seq_len = int(seq_len)
|
| self.pre_len = int(pre_len)
|
| self.scaler = None
|
| self.split = None
|
|
|
| def __getitem__(self, index):
|
| s_begin = index
|
| s_end = s_begin + self.seq_len
|
| r_end = s_end + self.pre_len
|
|
|
| x = self.split[s_begin:s_end]
|
| y = self.split[s_end:r_end]
|
| return x, y
|
|
|
| def __len__(self):
|
| if self.split is None:
|
| return 0
|
| return max(0, len(self.split) - self.seq_len - self.pre_len)
|
|
|
|
|
| class Dataset_Dhfm(_BaseTimeSeriesDataset):
|
| def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None):
|
| super().__init__(flag, seq_len, pre_len)
|
| self.path = root_path
|
|
|
| load_data = np.load(root_path)
|
| data = np.array(load_data).transpose()
|
| data = _to_float32(data)
|
|
|
| train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
|
| train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
|
|
|
| if self.flag == "train":
|
| self.split = train_data
|
| elif self.flag == "val":
|
| self.split = val_data
|
| else:
|
| self.split = test_data
|
|
|
|
|
| class Dataset_ECG(_BaseTimeSeriesDataset):
|
| def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None):
|
| super().__init__(flag, seq_len, pre_len)
|
| self.path = root_path
|
|
|
| df = pd.read_csv(root_path)
|
| data = _clean_numeric_csv(df)
|
|
|
| train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
|
| train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
|
|
|
| if self.flag == "train":
|
| self.split = train_data
|
| elif self.flag == "val":
|
| self.split = val_data
|
| else:
|
| self.split = test_data
|
|
|
| class Dataset_Solar(_BaseTimeSeriesDataset):
|
| def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None):
|
| super().__init__(flag, seq_len, pre_len)
|
| self.path = root_path
|
|
|
| files = os.listdir(root_path)
|
| solar_data = []
|
| time_data = None
|
|
|
| for file in files:
|
| full = os.path.join(root_path, file)
|
| if os.path.isdir(full):
|
| continue
|
| if file.startswith("DA_"):
|
| arr = pd.read_csv(full).values
|
| raw_time = arr[:, 0:1]
|
| if time_data is None:
|
| time_data = raw_time
|
| raw_data = arr[:, 1:arr.shape[1]]
|
| raw_data = raw_data.transpose()
|
| solar_data.append(raw_data)
|
|
|
| if len(solar_data) == 0 or time_data is None:
|
| raise ValueError(f"No solar files found in {root_path} with prefix 'DA_'.")
|
|
|
| solar_data = np.array(solar_data).squeeze(1).transpose()
|
| time_data = np.array(time_data)
|
| out = np.concatenate((time_data, solar_data), axis=1)
|
|
|
| filtered = []
|
| for item in out:
|
| dt = datetime.datetime.strptime(item[0], "%m/%d/%y %H:%M")
|
| if 8 <= dt.hour <= 17:
|
| filtered.append(item[1:out.shape[1]-1])
|
|
|
| data = _to_float32(np.array(filtered))
|
|
|
| train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
|
| train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
|
|
|
| if self.flag == "train":
|
| self.split = train_data
|
| elif self.flag == "val":
|
| self.split = val_data
|
| else:
|
| self.split = test_data
|
|
|
|
|
| class Dataset_Wiki(_BaseTimeSeriesDataset):
|
| def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None):
|
| super().__init__(flag, seq_len, pre_len)
|
| self.path = root_path
|
|
|
| df = pd.read_csv(root_path)
|
|
|
| if df.shape[1] < 2:
|
| raise ValueError("Wiki CSV must have at least 2 columns (time + features).")
|
|
|
| df_feat = df.iloc[:, 1:]
|
| data = _clean_numeric_csv(df_feat)
|
|
|
| train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
|
| train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
|
|
|
| if self.flag == "train":
|
| self.split = train_data
|
| elif self.flag == "val":
|
| self.split = val_data
|
| else:
|
| self.split = test_data
|
|
|
|
|
|
|
| class Dataset_PEMS_BAY(_BaseTimeSeriesDataset):
|
| def __init__(self, root_path, flag, seq_len, pre_len, type, train_ratio, val_ratio, scaler=None, fillna="ffill"):
|
| super().__init__(flag, seq_len, pre_len)
|
| self.path = root_path
|
|
|
| obj = pd.read_hdf(root_path)
|
|
|
| if isinstance(obj, pd.Series):
|
| df = obj.to_frame()
|
| elif isinstance(obj, pd.DataFrame):
|
| df = obj
|
| else:
|
| df = pd.DataFrame(obj)
|
|
|
| if fillna == "ffill":
|
| df = df.ffill()
|
| df = df.fillna(0.0)
|
| elif fillna == "zero":
|
| df = df.fillna(0.0)
|
| elif fillna == "drop":
|
| df = df.dropna(axis=0, how="any")
|
| elif fillna is None:
|
| pass
|
| else:
|
| raise ValueError("fillna must be one of: 'ffill', 'zero', 'drop', or None")
|
|
|
| data = df.values.astype(np.float32)
|
|
|
| train_data, val_data, test_data = _split_with_overlap(data, train_ratio, val_ratio, self.seq_len)
|
| train_data, val_data, test_data, self.scaler = _fit_transform_splits(train_data, val_data, test_data, type, scaler)
|
|
|
| if self.flag == "train":
|
| self.split = train_data
|
| elif self.flag == "val":
|
| self.split = val_data
|
| else:
|
| self.split = test_data
|
|
|
