gps full

examples/LennardJones/LJ.json

@@ -19,7 +19,9 @@
   "NeuralNetwork": {
     "Architecture": {
       "periodic_boundary_conditions": true,
-      "model_type": "DimeNet",
+      "global_attn_engine": "GPS",
+      "global_attn_type": "multihead",
+      "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": {

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):

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"][
@@ -271,7 +271,11 @@
         os.environ["HYDRAGNN_AGGR_BACKEND"] = "mpi"
         os.environ["HYDRAGNN_USE_ddstore"] = "1"
 
-    (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders(
+    (
+        train_loader,
+        val_loader,
+        test_loader,
+    ) = hydragnn.preprocess.create_dataloaders(
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 

examples/alexandria/alexandria_energy.json

@@ -4,7 +4,7 @@
   },
   "NeuralNetwork": {
     "Architecture": {
-      "model_type": "EGNN",
+      "mpnn_type": "EGNN",
       "equivariance": true,
       "radius": 5,
       "max_neighbours": 100000,

examples/alexandria/alexandria_forces.json

@@ -4,7 +4,7 @@
   },
   "NeuralNetwork": {
     "Architecture": {
-      "model_type": "EGNN",
+      "mpnn_type": "EGNN",
       "equivariance": true,
       "radius": 5,
       "max_neighbours": 100000,

examples/alexandria/train.py

@@ -530,7 +530,11 @@ def get(self, idx):
         os.environ["HYDRAGNN_AGGR_BACKEND"] = "mpi"
         os.environ["HYDRAGNN_USE_ddstore"] = "1"
 
-    (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders(
+    (
+        train_loader,
+        val_loader,
+        test_loader,
+    ) = hydragnn.preprocess.create_dataloaders(
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 

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,

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,

examples/ani1_x/train.py

@@ -389,7 +389,11 @@ def get(self, idx):
         os.environ["HYDRAGNN_AGGR_BACKEND"] = "mpi"
         os.environ["HYDRAGNN_USE_ddstore"] = "1"
 
-    (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders(
+    (
+        train_loader,
+        val_loader,
+        test_loader,
+    ) = hydragnn.preprocess.create_dataloaders(
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 

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,

examples/csce/train_gap.py

@@ -383,7 +383,11 @@ def __getitem__(self, idx):
         os.environ["HYDRAGNN_AGGR_BACKEND"] = "mpi"
         os.environ["HYDRAGNN_USE_ddstore"] = "1"
 
-    (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders(
+    (
+        train_loader,
+        val_loader,
+        test_loader,
+    ) = hydragnn.preprocess.create_dataloaders(
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
     comm.Barrier()

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,

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,

examples/dftb_uv_spectrum/train_discrete_uv_spectrum.py

@@ -317,7 +317,11 @@ def get(self, idx):
         os.environ["HYDRAGNN_AGGR_BACKEND"] = "mpi"
         os.environ["HYDRAGNN_USE_ddstore"] = "1"
 
-    (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders(
+    (
+        train_loader,
+        val_loader,
+        test_loader,
+    ) = hydragnn.preprocess.create_dataloaders(
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 

examples/dftb_uv_spectrum/train_smooth_uv_spectrum.py

@@ -315,7 +315,11 @@ def get(self, idx):
         os.environ["HYDRAGNN_AGGR_BACKEND"] = "mpi"
         os.environ["HYDRAGNN_USE_ddstore"] = "1"
 
-    (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders(
+    (
+        train_loader,
+        val_loader,
+        test_loader,
+    ) = hydragnn.preprocess.create_dataloaders(
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 

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"],

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"],

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"],

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"],

examples/eam/eam.py

@@ -167,7 +167,11 @@ def info(*args, logtype="info", sep=" "):
         % (len(trainset), len(valset), len(testset))
     )
 
-    (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders(
+    (
+        train_loader,
+        val_loader,
+        test_loader,
+    ) = hydragnn.preprocess.create_dataloaders(
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
     timer.stop()

examples/ising_model/create_configurations.py

@@ -40,7 +40,7 @@ def E_dimensionless(config, L, spin_function, scale_spin):
                 spin[x, y, z] = spin_function(config[x, y, z])
 
     count_pos = 0
-    number_nodes = L ** 3
+    number_nodes = L**3
     positions = np.zeros((number_nodes, 3))
     atomic_features = np.zeros((number_nodes, 5))
     for x in range(L):
@@ -79,15 +79,15 @@ def create_dataset(
 
     count_config = 0
 
-    for num_downs in tqdm(range(0, L ** 3)):
+    for num_downs in tqdm(range(0, L**3)):
 
-        primal_configuration = np.ones((L ** 3,))
+        primal_configuration = np.ones((L**3,))
         for down in range(0, num_downs):
             primal_configuration[down] = -1.0
 
         # If the current composition has a total number of possible configurations above
         # the hard cutoff threshold, a random configurational subset is picked
-        if scipy.special.binom(L ** 3, num_downs) > histogram_cutoff:
+        if scipy.special.binom(L**3, num_downs) > histogram_cutoff:
             for num_config in range(0, histogram_cutoff):
                 config = np.random.permutation(primal_configuration)
                 config = np.reshape(config, (L, L, L))

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,

examples/ising_model/train_ising.py

@@ -86,7 +86,7 @@ def create_dataset_mpi(
         np.random.seed(seed)
 
     count_config = 0
-    rx = list(nsplit(range(0, L ** 3), comm_size))[rank]
+    rx = list(nsplit(range(0, L**3), comm_size))[rank]
     info("rx", rx.start, rx.stop)
     for num_downs in range(rx.start, rx.stop):
         subdir = os.path.join(dir, str(num_downs))
@@ -98,13 +98,13 @@ def create_dataset_mpi(
         prefix = "output_%d_" % num_downs
         subdir = os.path.join(dir, str(num_downs))
 
-        primal_configuration = np.ones((L ** 3,))
+        primal_configuration = np.ones((L**3,))
         for down in range(0, num_downs):
             primal_configuration[down] = -1.0
 
         # If the current composition has a total number of possible configurations above
         # the hard cutoff threshold, a random configurational subset is picked
-        if scipy.special.binom(L ** 3, num_downs) > histogram_cutoff:
+        if scipy.special.binom(L**3, num_downs) > histogram_cutoff:
             for num_config in range(0, histogram_cutoff):
                 config = np.random.permutation(primal_configuration)
                 config = np.reshape(config, (L, L, L))
@@ -355,7 +355,11 @@ def info(*args, logtype="info", sep=" "):
         os.environ["HYDRAGNN_AGGR_BACKEND"] = "mpi"
         os.environ["HYDRAGNN_USE_ddstore"] = "1"
 
-    (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders(
+    (
+        train_loader,
+        val_loader,
+        test_loader,
+    ) = hydragnn.preprocess.create_dataloaders(
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
     timer.stop()

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,

examples/lsms/lsms.py

@@ -170,7 +170,11 @@ def info(*args, logtype="info", sep=" "):
         % (len(trainset), len(valset), len(testset))
     )
 
-    (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders(
+    (
+        train_loader,
+        val_loader,
+        test_loader,
+    ) = hydragnn.preprocess.create_dataloaders(
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
     timer.stop()

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": "CGCNN",
             "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,