| 193 | # https://doi.org/10.1016/j.trc.2023.104205 |
| 194 | class HSTGCN(nn.Module): |
| 195 | def __init__(self, seq, n_fea, adj_distance, adj_demand, alpha=0.5): |
| 196 | super(HSTGCN, self).__init__() |
| 197 | # hyper-params |
| 198 | self.nodes = adj_distance.shape[0] |
| 199 | self.alpha = alpha |
| 200 | hidden = seq - n_fea + 1 |
| 201 | |
| 202 | # network components |
| 203 | self.encoder = nn.Conv2d(self.nodes, self.nodes, (n_fea, n_fea)) |
| 204 | self.linear = nn.Linear(hidden, hidden) |
| 205 | self.distance_gcn_l1 = nn.Linear(hidden, hidden) |
| 206 | self.distance_gcn_l2 = nn.Linear(hidden, hidden) |
| 207 | self.gru1 = nn.GRU(self.nodes, self.nodes, num_layers=2, batch_first=True) |
| 208 | self.demand_gcn_l1 = nn.Linear(hidden, hidden) |
| 209 | self.demand_gcn_l2 = nn.Linear(hidden, hidden) |
| 210 | self.gru2 = nn.GRU(self.nodes, self.nodes, num_layers=2, batch_first=True) |
| 211 | self.decoder = nn.Sequential(nn.Linear(hidden, 16), |
| 212 | nn.ReLU(), |
| 213 | nn.Linear(16, 1) |
| 214 | ) |
| 215 | |
| 216 | self.act = nn.ReLU() |
| 217 | self.dropout = nn.Dropout(p=0.5) |
| 218 | |
| 219 | # calculate A_delta matrix |
| 220 | deg = torch.sum(adj_distance, dim=0) |
| 221 | deg = torch.diag(deg) |
| 222 | deg_delta = torch.linalg.inv(torch.sqrt(deg)) |
| 223 | a_delta = torch.matmul(torch.matmul(deg_delta, adj_distance), deg_delta) |
| 224 | self.A_dis = a_delta |
| 225 | |
| 226 | deg = torch.sum(adj_demand, dim=0) |
| 227 | deg = torch.diag(deg) |
| 228 | deg_delta = torch.linalg.inv(torch.sqrt(deg)) |
| 229 | a_delta = torch.matmul(torch.matmul(deg_delta, adj_demand), deg_delta) |
| 230 | self.A_dem = a_delta |
| 231 | |
| 232 | def forward(self, occ, prc): # occ.shape = [batch, node, seq] |
| 233 | x = torch.stack([occ, prc], dim=3) |