Implementing graphGPS as a global attention wrapper around Hydra MPNNs (

#315) * gps full * black formatted * updated tests * black formatted * Fix SchNet * additional tests added and pytest input fixed * deepspeed test added * removed edge_dim check * changed copyright year * SchNet threahsold increased for conv_head --------- Co-authored-by: Rylie Weaver <rylieweaver9@gmail.com> Co-authored-by: allaffa <massimiliano.lupo.pasini@gmail.com>
ORNL · Dec 13, 2024 · 1419dc4 · 1419dc4
1 parent 18251d0
commit 1419dc4
Show file tree

Hide file tree

Showing 76 changed files with 1,277 additions and 266 deletions.
diff --git a/examples/LennardJones/LJ.json b/examples/LennardJones/LJ.json
@@ -19,7 +19,9 @@
   "NeuralNetwork": {
     "Architecture": {
       "periodic_boundary_conditions": true,
-      "model_type": "DimeNet",
+      "global_attn_engine": "",
+      "global_attn_type": "",
+      "mpnn_type": "EGNN",
       "radius": 5.0,
       "max_neighbours": 5,
       "int_emb_size": 32,
@@ -33,8 +35,10 @@
       "max_ell": 1,
       "node_max_ell": 1,
       "num_radial": 5,
+      "pe_dim": 6,
+      "global_attn_heads": 8,
       "num_spherical": 2,
-      "hidden_dim": 20,
+      "hidden_dim": 32,
       "num_conv_layers": 4,
       "output_heads": {
         "node": {

diff --git a/examples/LennardJones/LJ_data.py b/examples/LennardJones/LJ_data.py
@@ -20,6 +20,7 @@
 # Torch
 import torch
 from torch_geometric.data import Data
+from torch_geometric.transforms import AddLaplacianEigenvectorPE
 
 # torch.set_default_tensor_type(torch.DoubleTensor)
 # torch.set_default_dtype(torch.float64)
@@ -101,11 +102,19 @@ def __init__(self, dirpath, config, dist=False, sampling=None):
 
         rx = list(nsplit((dirfiles), self.world_size))[self.rank]
 
+        # LPE
+        self.transform = AddLaplacianEigenvectorPE(
+            k=config["NeuralNetwork"]["Architecture"]["pe_dim"],
+            attr_name="pe",
+            is_undirected=True,
+        )
+
         for file in rx:
             filepath = os.path.join(dirpath, file)
             self.dataset.append(self.transform_input_to_data_object_base(filepath))
 
     def transform_input_to_data_object_base(self, filepath):
+
         # Using readline()
         file = open(filepath, "r")
 
@@ -183,7 +192,13 @@ def transform_input_to_data_object_base(self, filepath):
         # Create pbc edges and lengths
         edge_creation = get_radius_graph_pbc(self.radius, self.max_neighbours)
         data = edge_creation(data)
-
+        data = self.transform(data)
+        # gps requires relative edge features, introduced rel_lapPe as edge encodings
+        source_pe = data.pe[data.edge_index[0]]
+        target_pe = data.pe[data.edge_index[1]]
+        data.rel_pe = torch.abs(
+            source_pe - target_pe
+        )  # Compute feature-wise difference
         return data
 
     def len(self):

diff --git a/examples/LennardJones/LennardJones.py b/examples/LennardJones/LennardJones.py
@@ -64,7 +64,7 @@
         help="preprocess only (no training)",
     )
     parser.add_argument("--inputfile", help="input file", type=str, default="LJ.json")
-    parser.add_argument("--model_type", help="model type", type=str, default=None)
+    parser.add_argument("--mpnn_type", help="model type", type=str, default=None)
     parser.add_argument("--mae", action="store_true", help="do mae calculation")
     parser.add_argument("--ddstore", action="store_true", help="ddstore dataset")
     parser.add_argument("--ddstore_width", type=int, help="ddstore width", default=None)
@@ -102,10 +102,10 @@
     # Configurable run choices (JSON file that accompanies this example script).
     with open(input_filename, "r") as f:
         config = json.load(f)
-    config["NeuralNetwork"]["Architecture"]["model_type"] = (
-        args.model_type
-        if args.model_type
-        else config["NeuralNetwork"]["Architecture"]["model_type"]
+    config["NeuralNetwork"]["Architecture"]["mpnn_type"] = (
+        args.mpnn_type
+        if args.mpnn_type
+        else config["NeuralNetwork"]["Architecture"]["mpnn_type"]
     )
     verbosity = config["Verbosity"]["level"]
     config["NeuralNetwork"]["Variables_of_interest"][

diff --git a/examples/alexandria/alexandria_energy.json b/examples/alexandria/alexandria_energy.json
@@ -4,7 +4,7 @@
   },
   "NeuralNetwork": {
     "Architecture": {
-      "model_type": "EGNN",
+      "mpnn_type": "EGNN",
       "equivariance": true,
       "radius": 5,
       "max_neighbours": 100000,

diff --git a/examples/alexandria/alexandria_forces.json b/examples/alexandria/alexandria_forces.json
@@ -4,7 +4,7 @@
   },
   "NeuralNetwork": {
     "Architecture": {
-      "model_type": "EGNN",
+      "mpnn_type": "EGNN",
       "equivariance": true,
       "radius": 5,
       "max_neighbours": 100000,

diff --git a/examples/ani1_x/ani1x_energy.json b/examples/ani1_x/ani1x_energy.json
@@ -4,7 +4,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "EGNN",
+            "mpnn_type": "EGNN",
             "equivariance": true,
             "radius": 5.0,
             "max_neighbours": 100000,

diff --git a/examples/ani1_x/ani1x_forces.json b/examples/ani1_x/ani1x_forces.json
@@ -4,7 +4,7 @@
     },
    "NeuralNetwork": {
         "Architecture": {
-            "model_type": "EGNN",
+            "mpnn_type": "EGNN",
             "equivariance": true,
             "radius": 5.0,
             "max_neighbours": 100000,

diff --git a/examples/csce/csce_gap.json b/examples/csce/csce_gap.json
@@ -4,7 +4,7 @@
   },
   "NeuralNetwork": {
     "Architecture": {
-      "model_type": "PNA",
+      "mpnn_type": "PNA",
       "max_neighbours": 20,
       "hidden_dim": 200,
       "num_conv_layers": 6,

diff --git a/examples/dftb_uv_spectrum/dftb_discrete_uv_spectrum.json b/examples/dftb_uv_spectrum/dftb_discrete_uv_spectrum.json
@@ -4,7 +4,7 @@
   },
   "NeuralNetwork": {
     "Architecture": {
-      "model_type": "PNA",
+      "mpnn_type": "PNA",
       "max_neighbours": 20,
       "hidden_dim": 200,
       "num_conv_layers": 6,

diff --git a/examples/dftb_uv_spectrum/dftb_smooth_uv_spectrum.json b/examples/dftb_uv_spectrum/dftb_smooth_uv_spectrum.json
@@ -4,7 +4,7 @@
   },
   "NeuralNetwork": {
     "Architecture": {
-      "model_type": "PNA",
+      "mpnn_type": "PNA",
       "max_neighbours": 20,
       "hidden_dim": 200,
       "num_conv_layers": 6,

diff --git a/examples/eam/NiNb_EAM_bulk.json b/examples/eam/NiNb_EAM_bulk.json
@@ -21,7 +21,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "PNA",
+            "mpnn_type": "PNA",
             "radius": 3.0,
             "max_neighbours": 100000,
 	    "edge_features": ["lengths"],

diff --git a/examples/eam/NiNb_EAM_bulk_multitask.json b/examples/eam/NiNb_EAM_bulk_multitask.json
@@ -21,7 +21,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "PNA",
+            "mpnn_type": "PNA",
             "radius": 3,
             "max_neighbours": 100000,
 	    "edge_features": ["lengths"],

diff --git a/examples/eam/NiNb_EAM_energy.json b/examples/eam/NiNb_EAM_energy.json
@@ -21,7 +21,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "PNA",
+            "mpnn_type": "PNA",
             "radius": 3.0,
             "max_neighbours": 100000,
 	    "edge_features": ["lengths"],

diff --git a/examples/eam/NiNb_EAM_multitask.json b/examples/eam/NiNb_EAM_multitask.json
@@ -21,7 +21,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "PNA",
+            "mpnn_type": "PNA",
             "radius": 3,
             "max_neighbours": 100000,
 	    "edge_features": ["lengths"],

diff --git a/examples/ising_model/ising_model.json b/examples/ising_model/ising_model.json
@@ -21,7 +21,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "PNA",
+            "mpnn_type": "PNA",
             "radius": 7,
             "max_neighbours": 100000,
             "hidden_dim": 20,

diff --git a/examples/lsms/lsms.json b/examples/lsms/lsms.json
@@ -21,7 +21,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "PNA",
+            "mpnn_type": "PNA",
             "radius": 7,
             "max_neighbours": 100,
             "periodic_boundary_conditions": false,

diff --git a/examples/md17/md17.json b/examples/md17/md17.json
@@ -5,11 +5,15 @@
     "NeuralNetwork": {
         "Profile": {"enable": 1},
         "Architecture": {
-            "model_type": "GIN",
+            "global_attn_engine": "GPS",
+            "global_attn_type": "multihead",
+            "mpnn_type": "SchNet",
             "radius": 7,
             "max_neighbours": 5,
             "periodic_boundary_conditions": false,
-            "hidden_dim": 5,
+            "pe_dim": 6,
+            "global_attn_heads": 8,
+            "hidden_dim": 64,
             "num_conv_layers": 6,
             "int_emb_size": 32,
             "out_emb_size": 16,

diff --git a/examples/md17/md17.py b/examples/md17/md17.py
@@ -1,7 +1,9 @@
 import os
+import pdb
 import json
 import torch
 import torch_geometric
+from torch_geometric.transforms import AddLaplacianEigenvectorPE
 import argparse
 
 # deprecated in torch_geometric 2.0
@@ -14,14 +16,19 @@
 
 
 # Update each sample prior to loading.
-def md17_pre_transform(data, compute_edges):
+def md17_pre_transform(data, compute_edges, transform):
     # Set descriptor as element type.
     data.x = data.z.float().view(-1, 1)
     # Only predict energy (index 0 of 2 properties) for this run.
     data.y = data.energy / len(data.x)
     graph_features_dim = [1]
     node_feature_dim = [1]
     data = compute_edges(data)
+    data = transform(data)
+    # gps requires relative edge features, introduced rel_lapPe as edge encodings
+    source_pe = data.pe[data.edge_index[0]]
+    target_pe = data.pe[data.edge_index[1]]
+    data.rel_pe = torch.abs(source_pe - target_pe)  # Compute feature-wise difference
     return data
 
 
@@ -30,7 +37,7 @@ def md17_pre_filter(data):
     return torch.rand(1) < 0.25
 
 
-def main(model_type=None):
+def main(mpnn_type=None, global_attn_engine=None, global_attn_type=None):
     # FIX random seed
     random_state = 0
     torch.manual_seed(random_state)
@@ -47,26 +54,41 @@ def main(model_type=None):
     arch_config = config["NeuralNetwork"]["Architecture"]
 
     # If a model type is provided, update the configuration
-    if model_type:
-        config["NeuralNetwork"]["Architecture"]["model_type"] = model_type
+    if global_attn_engine:
+        config["NeuralNetwork"]["Architecture"][
+            "global_attn_engine"
+        ] = global_attn_engine
+
+    if global_attn_type:
+        config["NeuralNetwork"]["Architecture"]["global_attn_type"] = global_attn_type
+
+    if mpnn_type:
+        config["NeuralNetwork"]["Architecture"]["mpnn_type"] = mpnn_type
 
     # Always initialize for multi-rank training.
     world_size, world_rank = hydragnn.utils.distributed.setup_ddp()
 
-    log_name = f"md17_test_{model_type}" if model_type else "md17_test"
+    log_name = f"md17_test_{mpnn_type}" if mpnn_type else "md17_test"
     # Enable print to log file.
     hydragnn.utils.print.print_utils.setup_log(log_name)
 
     # Preprocess configurations for edge computation
     compute_edges = hydragnn.preprocess.get_radius_graph_config(arch_config)
 
+    # LPE
+    transform = AddLaplacianEigenvectorPE(
+        k=config["NeuralNetwork"]["Architecture"]["pe_dim"],
+        attr_name="pe",
+        is_undirected=True,
+    )
+
     # Fix for MD17 datasets
     torch_geometric.datasets.MD17.file_names["uracil"] = "md17_uracil.npz"
 
     dataset = torch_geometric.datasets.MD17(
         root="dataset/md17",
         name="uracil",
-        pre_transform=lambda data: md17_pre_transform(data, compute_edges),
+        pre_transform=lambda data: md17_pre_transform(data, compute_edges, transform),
         pre_filter=md17_pre_filter,
     )
     train, val, test = hydragnn.preprocess.split_dataset(
@@ -115,10 +137,25 @@ def main(model_type=None):
         description="Run MD17 example with an optional model type."
     )
     parser.add_argument(
-        "--model_type",
+        "--mpnn_type",
         type=str,
         default=None,
         help="Specify the model type for training (default: None).",
     )
+
+    parser.add_argument(
+        "--global_attn_engine",
+        type=str,
+        default=None,
+        help="Specify if global attention is being used (default: None).",
+    )
+
+    parser.add_argument(
+        "--global_attn_type",
+        type=str,
+        default=None,
+        help="Specify the global attention type (default: None).",
+    )
+
     args = parser.parse_args()
-    main(model_type=args.model_type)
+    main(mpnn_type=args.mpnn_type)
diff --git a/examples/mptrj/mptrj_energy.json b/examples/mptrj/mptrj_energy.json
@@ -4,7 +4,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "EGNN",
+            "mpnn_type": "EGNN",
             "equivariance": true,
             "radius": 5.0,
             "max_neighbours": 100000,

diff --git a/examples/mptrj/mptrj_forces.json b/examples/mptrj/mptrj_forces.json
@@ -4,7 +4,7 @@
     },
    "NeuralNetwork": {
         "Architecture": {
-            "model_type": "EGNN",
+            "mpnn_type": "EGNN",
             "equivariance": true,
             "radius": 5.0,
             "max_neighbours": 100000,

diff --git a/examples/multidataset/gfm_energy.json b/examples/multidataset/gfm_energy.json
@@ -4,7 +4,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "EGNN",
+            "mpnn_type": "EGNN",
             "equivariance": true,
             "radius": 5.0,
             "max_neighbours": 100000,

diff --git a/examples/multidataset/gfm_forces.json b/examples/multidataset/gfm_forces.json
@@ -4,7 +4,7 @@
     },
    "NeuralNetwork": {
         "Architecture": {
-            "model_type": "EGNN",
+            "mpnn_type": "EGNN",
             "equivariance": true,
             "radius": 5.0,
             "max_neighbours": 100000,

diff --git a/examples/multidataset/gfm_multitasking.json b/examples/multidataset/gfm_multitasking.json
@@ -4,7 +4,7 @@
     },
    "NeuralNetwork": {
         "Architecture": {
-            "model_type": "EGNN",
+            "mpnn_type": "EGNN",
             "equivariance": true,
             "radius": 5.0,
             "max_neighbours": 100000,

diff --git a/examples/multidataset_hpo/gfm_energy.json b/examples/multidataset_hpo/gfm_energy.json
@@ -4,7 +4,7 @@
     },
     "NeuralNetwork": {
         "Architecture": {
-            "model_type": "EGNN",
+            "mpnn_type": "EGNN",
             "equivariance": true,
             "radius": 5.0,
             "max_neighbours": 100000,