| 407 | |
| 408 | |
| 409 | def forward(self, batch: HeteroData): |
| 410 | graph_x_dict = batch.x_dict |
| 411 | |
| 412 | # vitual global node |
| 413 | pq_x = torch.stack(torch.chunk(graph_x_dict['PQ'], self.batch_size, dim=0), dim=0) # B, 29, D |
| 414 | pv_x = torch.stack(torch.chunk(graph_x_dict['PV'], self.batch_size, dim=0), dim=0) |
| 415 | slack_x = torch.stack(torch.chunk(graph_x_dict['Slack'], self.batch_size, dim=0), dim=0) |
| 416 | global_feature = torch.cat((pq_x,pv_x,slack_x), dim=1) # B, (29+9+1), D |
| 417 | global_feature = self.global_transform(global_feature) |
| 418 | global_feature_mean = global_feature.mean(dim=1, keepdim=True) |
| 419 | global_feature_max, _ = global_feature.max(dim=1, keepdim=True) |
| 420 | |
| 421 | # forward gcn |
| 422 | for layer in self.graph_layers: |
| 423 | graph_x_dict = layer(graph_x_dict, |
| 424 | batch.edge_index_dict, |
| 425 | batch.edge_attr_dict) |
| 426 | ## NEW: add non-linear |
| 427 | graph_x_dict = {key: self.act_layer(x) for key, x in graph_x_dict.items()} |
| 428 | |
| 429 | global_node_feat = torch.cat([global_feature_mean, global_feature_max], dim=1) |
| 430 | |
| 431 | # cross attent the global feat. |
| 432 | res = {} |
| 433 | for key in ["PQ", "PV"]: |
| 434 | # get size |
| 435 | BN, K = batch[key].x.size() |
| 436 | B = self.batch_size |
| 437 | N = BN // B |
| 438 | # ca |
| 439 | graph_x_dict[key] = graph_x_dict[key] + self.cross_attention(graph_x_dict[key].view(B, N, K), global_node_feat).contiguous().view(B*N, K) |
| 440 | # norm |
| 441 | res[key] = self.norm(graph_x_dict[key]) |
| 442 | res["Slack"] = graph_x_dict["Slack"] |
| 443 | |
| 444 | return res |
| 445 | |
| 446 | |
| 447 | # ----- ffn layers |