Incorporate Graphormer

examples/config/graphormer_pretraining.yaml

@@ -0,0 +1,56 @@
+callbacks:
+  patience: 100
+  tol: 0
+data:
+  baseMVA: 100
+  learn_mask: false
+  mask_dim: 6
+  mask_ratio: 0.5
+  mask_type: rnd
+  mask_value: -1.0
+  networks:
+  # - Texas2k_case1_2016summerpeak
+  - case24_ieee_rts
+  - case118_ieee
+  - case300_ieee
+  normalization: baseMVAnorm
+  scenarios:
+  - 5000
+  - 5000
+  - 5000
+  test_ratio: 0.1
+  val_ratio: 0.1
+  workers: 4
+  add_graphormer_encoding: true
+  max_node_num: 300 # necessary for Graphormer
+  max_hops: 6 # for the edge encoding, should match 
+  edge_type: multi_hop # singlehop
+model:
+  attention_head: 8
+  dropout: 0.1
+  edge_dim: 2
+  hidden_size: 123
+  input_dim: 9
+  num_layers: 14
+  output_dim: 6
+  pe_dim: 20
+  type: Graphormer  #GPSTransformer  #
+optimizer:
+  beta1: 0.9
+  beta2: 0.999
+  learning_rate: 0.0001
+  lr_decay: 0.7
+  lr_patience: 10
+seed: 0
+training:
+  batch_size: 8
+  epochs: 500
+  loss_weights:
+  - 0.01
+  - 0.99
+  losses:
+  - MaskedMSE
+  - PBE
+  accelerator: auto
+  devices: auto
+  strategy: auto

gridfm_graphkit/datasets/powergrid_datamodule.py

@@ -128,6 +128,7 @@ def setup(self, stage: str):
                 pe_dim=self.args.model.pe_dim,
                 mask_dim=self.args.data.mask_dim,
                 transform=get_transform(args=self.args),
+                args=self.args.data
             )
             self.datasets.append(dataset)
 

gridfm_graphkit/datasets/powergrid_dataset.py

@@ -2,6 +2,7 @@
 from gridfm_graphkit.datasets.transforms import (
     AddEdgeWeights,
     AddNormalizedRandomWalkPE,
+    AddGraphormerEncodings
 )
 
 import os.path as osp
@@ -43,6 +44,7 @@ def __init__(
         transform: Optional[Callable] = None,
         pre_transform: Optional[Callable] = None,
         pre_filter: Optional[Callable] = None,
+        args: Optional = {},
     ):
         self.norm_method = norm_method
         self.node_normalizer = node_normalizer
@@ -51,6 +53,13 @@ def __init__(
         self.mask_dim = mask_dim
         self.length = None
 
+        if ("add_graphormer_encoding" in args) and args.add_graphormer_encoding:
+            self.add_graphormer_encoding = args.add_graphormer_encoding
+            self.max_node_num = args.max_node_num
+            self.max_hops = args.max_hops
+        else:
+            self.add_graphormer_encoding = False
+
         super().__init__(root, transform, pre_transform, pre_filter)
 
         # Load normalization stats if available
@@ -171,6 +180,7 @@ def process(self):
                 attr_name="pe",
             )
             graph_data = pe_transform(graph_data)
+
             torch.save(
                 graph_data,
                 osp.join(
@@ -194,6 +204,11 @@ def len(self):
             self.length = len(files)
         return self.length
 
+    def __cat_dim__(self, key, value, *args, **kwargs):
+        if key in ['attn_bias', 'spatial_pos', 'in_degree', 'edge_input']:
+            return None
+        return super().__cat_dim__(key, value, *args, **kwargs)
+
     def get(self, idx):
         file_name = osp.join(
             self.processed_dir,
@@ -204,6 +219,14 @@ def get(self, idx):
         data = torch.load(file_name, weights_only=False)
         if self.transform:
             data = self.transform(data)
+
+        if self.add_graphormer_encoding:
+            gr_transform = AddGraphormerEncodings(
+                                        self.max_node_num,
+                                        self.max_hops
+                                        )
+            data = gr_transform(data)
+
         return data
 
     def change_transform(self, new_transform):

gridfm_graphkit/datasets/transforms.py

@@ -4,7 +4,7 @@
 import torch
 from torch import Tensor
 from torch_geometric.transforms import BaseTransform
-from typing import Optional
+from typing import Optional, Any
 import torch_geometric.typing
 from torch_geometric.data import Data
 from torch_geometric.utils import (
@@ -15,6 +15,12 @@
     to_torch_csr_tensor,
 )
 
+import numpy as np
+import os
+import pyximport
+pyximport.install(setup_args={'include_dirs': np.get_include()})
+import gridfm_graphkit.models.algos as algos
+
 
 class AddNormalizedRandomWalkPE(BaseTransform):
     r"""Adds the random walk positional encoding from the
@@ -84,6 +90,184 @@ def get_pe(out: Tensor) -> Tensor:
         return data
 
 
+def add_node_attr(data: Data, 
+                    value: Any,
+                    attr_name: str
+                    ) -> Data:
+    if attr_name is None:
+        raise ValueError("Expected attr_name to be not None")
+    else:
+        data[attr_name] = value
+
+    return data
+
+def convert_to_single_emb(x, offset=512):
+    """
+    The edge feature embedding range is set to 512, with the futher assumption
+    that the input range is from -512 to 512. This may need to change in the future.
+    """
+    x = torch.clamp(x, -512, 512)
+    x = ( 512*(x+512)/1024 ).long()
+
+    feature_num = x.size(1) if len(x.size()) > 1 else 1
+    feature_offset = 1 + \
+        torch.arange(
+                0, 
+                (feature_num) * offset, 
+                offset,
+                 dtype=torch.long
+                 )
+
+    x = x + feature_offset
+
+    return x
+
+def get_edge_encoding(edge_attr, N, edge_index, max_dist, path):
+    if len(edge_attr.size()) == 1:
+            edge_attr = edge_attr[:, None]
+    attn_edge_type = torch.zeros([N, N, edge_attr.size(-1)], dtype=torch.long)
+    attn_edge_type[edge_index[0, :], edge_index[1, :]
+                    ] = convert_to_single_emb(edge_attr.long()) + 1
+    if os.name == 'nt':
+        edge_input = algos.gen_edge_input(
+                    max_dist, 
+                    path, 
+                    attn_edge_type.numpy(),
+                    localtype=np.int32
+                    )
+    else:
+        edge_input = algos.gen_edge_input(max_dist, path, attn_edge_type.numpy())
+
+    return attn_edge_type, torch.from_numpy(edge_input).long()
+
+def preprocess_item(data, max_hops):
+    """
+    Calculation of the Graphormer attention bias, and positional/structural 
+    variables. From a Data-like object the shortest paths in number of hops 
+    between nodes are calculated, being cut off at max_hops. In addition to the
+    centrality (assume undirected graphs) and attention bias, these are the 
+    inputs to the model structural and positional encodings.
+    """
+    edge_index = data.edge_index
+    edge_attr = data.edge_attr
+    N = data.num_nodes
+    edge_adj = torch.sparse_coo_tensor(
+                                    edge_index,
+                                    torch.ones(edge_index.shape[1]).to(data.x.device),
+                                    [N, N]
+                                    )
+
+    adj = edge_adj.to_dense().to(torch.int16)
+
+    # get shortest paths in number of hops (shortest_path_result) and intermediate nodes
+    # for those shortest paths (path)
+    if os.name == 'nt':
+        shortest_path_result, path = algos.floyd_warshall(
+                        adj.numpy().astype(np.int32), 
+                        max_hops, 
+                        localtype=np.int32
+                        )
+    else:
+        shortest_path_result, path = algos.floyd_warshall(
+                        adj.numpy().astype(np.int32), 
+                        max_hops
+                        )
+
+    spatial_pos = torch.from_numpy((shortest_path_result)).long().to(data.x.device)
+    attn_bias = torch.zeros([N, N], dtype=torch.float).to(data.x.device) 
+
+    if edge_attr is not None:
+        attn_edge_type, edge_input = get_edge_encoding(edge_attr, N, edge_index, max_hops, path)
+    else:
+        edge_input = None
+        attn_edge_type = None
+
+    in_degree = adj.long().sum(dim=1).view(-1)
+    out_degree = adj.long().sum(dim=0).view(-1)
+    return attn_bias, spatial_pos, in_degree, out_degree, attn_edge_type, edge_input
+
+def pad_1d_unsqueeze(x, padlen):
+    xlen = x.size(0)
+    if xlen < padlen:
+        new_x = x.new_zeros([padlen], dtype=x.dtype)
+        new_x[:xlen] = x
+        x = new_x
+    return x.unsqueeze(0)
+
+def pad_attn_bias_unsqueeze(x, padlen):
+    xlen = x.size(0)
+    if xlen < padlen:
+        new_x = x.new_zeros(
+            [padlen, padlen], dtype=x.dtype).fill_(float('-inf'))   
+        new_x[:xlen, :xlen] = x
+        new_x[xlen:, :xlen] = 0     
+        x = new_x
+    return x.unsqueeze(0)
+
+def pad_edge_bias_unsqueeze(x, padlen):
+    xlen = x.size(0)
+    if xlen < padlen:
+        new_x = x.new_zeros(
+            (padlen, padlen) + x.size()[2:], dtype=x.dtype).fill_(int(0))
+        new_x[:xlen, :xlen] = x
+        new_x[xlen:, :xlen] = 0
+        x = new_x
+    return x.unsqueeze(0)
+
+def pad_spatial_pos_unsqueeze(x, padlen):
+    xlen = x.size(0)
+    if xlen < padlen:
+        new_x = x.new_zeros([padlen, padlen], dtype=x.dtype)
+        new_x[:xlen, :xlen] = x
+        x = new_x
+    return x.unsqueeze(0)
+
+
+class AddGraphormerEncodings(BaseTransform):
+    """Adds a positional encoding (node centrallity) to the given graph, as 
+    well as the attention and edge biases, as described in: Do transformers 
+    really perform badly for graph representation?, C. Ying et al., 2021.
+
+    Args:
+        max_node_num (int): The number of nodes in the largest graph considered.
+        max_hops (int): The maximum path length between nodes to consider for
+                        the edge encodings.
+    """
+
+    def __init__(
+        self,
+        max_node_num: int,
+        max_hops: int,
+    ) -> None:
+        self.max_node_num = max_node_num
+        self.max_hops = max_hops
+
+    def forward(self, data: Data) -> Data:
+        if data.edge_index is None:
+            raise ValueError("Expected data.edge_index to be not None")
+
+        N = data.num_nodes
+        if N is None:
+            raise ValueError("Expected data.num_nodes to be not None")
+
+        attn_bias, spatial_pos, in_degree, out_degree, attn_edge_type, edge_input = \
+                            preprocess_item(data, self.max_hops)
+
+        attn_bias = pad_attn_bias_unsqueeze(attn_bias, self.max_node_num)
+        spatial_pos = pad_spatial_pos_unsqueeze(spatial_pos, self.max_node_num)
+        in_degree = pad_1d_unsqueeze(in_degree, self.max_node_num).squeeze()
+        edge_input = pad_edge_bias_unsqueeze(edge_input, self.max_node_num)
+        attn_edge_type = pad_edge_bias_unsqueeze(attn_edge_type, self.max_node_num)
+
+        data = add_node_attr(data, attn_bias, attr_name='attn_bias')
+        data = add_node_attr(data, spatial_pos, attr_name='spatial_pos')
+        data = add_node_attr(data, in_degree, attr_name='in_degree')
+        data = add_node_attr(data, edge_input, attr_name='edge_input')
+        data = add_node_attr(data, attn_edge_type, attr_name='attn_edge_type')
+
+        return data
+
+
 class AddEdgeWeights(BaseTransform):
     """
     Computes and adds edge weight as the magnitude of complex admittance.

gridfm_graphkit/models/__init__.py

@@ -1,4 +1,5 @@
 from gridfm_graphkit.models.gps_transformer import GPSTransformer
 from gridfm_graphkit.models.gnn_transformer import GNN_TransformerConv
+from gridfm_graphkit.models.graphormer import Graphormer
 
-__all__ = ["GPSTransformer", "GNN_TransformerConv"]
+__all__ = ["GPSTransformer", "GNN_TransformerConv", "Graphormer"]