| 7 | from utils.timefeatures import time_features |
| 8 | |
| 9 | class LargeGraphTemporalLoader(object): |
| 10 | |
| 11 | def __init__( self, data_path, edge_path=None, lags = 3, p_len=1, partition = 8, train=True,timeenc=0): |
| 12 | |
| 13 | if 'C2TM' in data_path: |
| 14 | df = pd.read_csv(data_path)/(8*1024*1024) |
| 15 | df = df.drop('time',axis=1) |
| 16 | self._data = df.values |
| 17 | elif 'CBS' in data_path: |
| 18 | df = pd.read_csv(data_path) |
| 19 | df = df.drop('time',axis=1) |
| 20 | self._data = df.values |
| 21 | else: |
| 22 | self._data = np.load(data_path)[...,2:4]/10 |
| 23 | print(self._data.shape, 'Call ...') |
| 24 | |
| 25 | # self.timeenc = timeenc |
| 26 | # self.data_stamp = self._get_stamp(df[['time']]) |
| 27 | |
| 28 | self.edge_path = edge_path |
| 29 | self.lags = lags |
| 30 | self.p_len = p_len |
| 31 | self.partition = partition |
| 32 | self.train = train |
| 33 | |
| 34 | |
| 35 | |
| 36 | self.n_node = self._data.shape[1] |
| 37 | self.times = self._data.shape[0] - lags - p_len + 1 |
| 38 | print('data shape: ',self._data.shape, self.n_node) |
| 39 | |
| 40 | if edge_path is not None: |
| 41 | with open(edge_path, 'rb') as f: |
| 42 | self._adj_mx = pickle.load(f)["adj_mx"] |
| 43 | else: |
| 44 | self._adj_mx = np.diag([1 for _ in range(self._data.shape[1])]) |
| 45 | print(self._adj_mx.shape, 'adj_max shape ...') |
| 46 | |
| 47 | def _get_stamp(self, df_stamp): |
| 48 | |
| 49 | df_stamp['date'] = pd.to_datetime(df_stamp.time) |
| 50 | if self.timeenc == 0: |
| 51 | df_stamp['month'] = df_stamp.date.apply(lambda row: row.month, 1) |
| 52 | df_stamp['day'] = df_stamp.date.apply(lambda row: row.day, 1) |
| 53 | df_stamp['weekday'] = df_stamp.date.apply(lambda row: row.weekday(), 1) |
| 54 | df_stamp['hour'] = df_stamp.date.apply(lambda row: row.hour, 1) |
| 55 | data_stamp = df_stamp.drop(['date'], 1).values |
| 56 | elif self.timeenc == 1: |
| 57 | data_stamp = time_features(pd.to_datetime(df_stamp['time'].values), freq=self.freq) |
| 58 | data_stamp = data_stamp.transpose(1, 0) |
| 59 | |
| 60 | return data_stamp |
| 61 | |
| 62 | |
| 63 | def _partition_edges(self,idxes): |
| 64 | edge_rows = self._adj_mx[idxes] |
| 65 | edge = edge_rows[:, idxes] |
| 66 | |