Open catalyst 2022 (#214)

* open catalyst 2022 example added * turned off implicit parallelization in ase.io.read * black formatting * using only bond length as edge feature to enable use of equivariant features --------- Co-authored-by: Massimiliano Lupo Pasini <[email protected]> Co-authored-by: Massimiliano Lupo Pasini <[email protected]>
ORNL · Mar 26, 2024 · 284f9ef · 284f9ef
1 parent ca8e656
commit 284f9ef
Show file tree

Hide file tree

Showing 7 changed files with 553 additions and 10 deletions.
diff --git a/examples/open_catalyst_2020/open_catalyst_energy.json b/examples/open_catalyst_2020/open_catalyst_energy.json
@@ -18,7 +18,7 @@
             "num_radial": 6,
             "num_spherical": 7,
             "num_filters": 126,
-            "edge_features": ["coord_x", "coord_y", "coord_z"],
+            "edge_features": ["length"],
             "hidden_dim": 50,
             "num_conv_layers": 3,
             "output_heads": {

diff --git a/examples/open_catalyst_2020/open_catalyst_forces.json b/examples/open_catalyst_2020/open_catalyst_forces.json
@@ -18,7 +18,7 @@
             "num_radial": 6,
             "num_spherical": 7,
             "num_filters": 126,
-            "edge_features": ["coord_x", "coord_y", "coord_z"],
+            "edge_features": ["length"],
             "hidden_dim": 50,
             "num_conv_layers": 3,
             "output_heads": {

diff --git a/examples/open_catalyst_2020/train.py b/examples/open_catalyst_2020/train.py
@@ -13,8 +13,6 @@
 from torch import tensor
 from torch_geometric.data import Data
 
-from torch_geometric.transforms import Spherical, LocalCartesian
-
 import hydragnn
 from hydragnn.utils.time_utils import Timer
 from hydragnn.utils.model import print_model
@@ -44,10 +42,6 @@ def info(*args, logtype="info", sep=" "):
     getattr(logging, logtype)(sep.join(map(str, args)))
 
 
-# transform_coordinates = Spherical(norm=False, cat=False)
-transform_coordinates = LocalCartesian(norm=False, cat=False)
-
-
 class OpenCatalystDataset(AbstractBaseDataset):
     def __init__(self, dirpath, var_config, data_type, dist=False):
         super().__init__()

diff --git a/examples/open_catalyst_2020/utils/atoms_to_graphs.py b/examples/open_catalyst_2020/utils/atoms_to_graphs.py
@@ -13,12 +13,13 @@
 
 import torch
 from torch_geometric.data import Data
-from torch_geometric.transforms import Spherical, LocalCartesian
+from torch_geometric.transforms import Distance, Spherical, LocalCartesian
 
 from hydragnn.preprocess.utils import RadiusGraph, RadiusGraphPBC
 
 # transform_coordinates = Spherical(norm=False, cat=False)
-transform_coordinates = LocalCartesian(norm=False, cat=False)
+# transform_coordinates = LocalCartesian(norm=False, cat=False)
+transform_coordinates = Distance(norm=False, cat=False)
 
 
 class AtomsToGraphs:

diff --git a/examples/open_catalyst_2022/open_catalyst_energy.json b/examples/open_catalyst_2022/open_catalyst_energy.json
@@ -0,0 +1,58 @@
+{
+    "Verbosity": {
+        "level": 2
+    },
+    "NeuralNetwork": {
+        "Architecture": {
+            "model_type": "EGNN",
+            "equivariance": true,
+            "radius": 5.0,
+            "max_neighbours": 100000,
+            "num_gaussians": 50,
+            "envelope_exponent": 5,
+            "int_emb_size": 64,
+            "basis_emb_size": 8,
+            "out_emb_size": 128,
+            "num_after_skip": 2,
+            "num_before_skip": 1,
+            "num_radial": 6,
+            "num_spherical": 7,
+            "num_filters": 126,
+            "edge_features": ["length"],
+            "hidden_dim": 50,
+            "num_conv_layers": 3,
+            "output_heads": {
+                "graph":{
+                    "num_sharedlayers": 2,
+                    "dim_sharedlayers": 50,
+                    "num_headlayers": 2,
+                    "dim_headlayers": [50,25]
+                }
+            },
+            "task_weights": [1.0]
+        },
+        "Variables_of_interest": {
+            "input_node_features": [0, 1, 2, 3],
+            "output_names": ["energy"],
+            "output_index": [0],
+            "output_dim": [1],
+            "type": ["graph"]
+        },
+        "Training": {
+            "num_epoch": 50,
+            "perc_train": 0.8,
+            "loss_function_type": "mae",
+            "batch_size": 32,
+            "continue": 0,
+            "Optimizer": {
+                "type": "AdamW",
+                "learning_rate": 1e-3
+            }
+        }
+    },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
+    }
+}
diff --git a/examples/open_catalyst_2022/open_catalyst_forces.json b/examples/open_catalyst_2022/open_catalyst_forces.json
@@ -0,0 +1,58 @@
+{
+    "Verbosity": {
+        "level": 2
+    },
+   "NeuralNetwork": {
+        "Architecture": {
+            "model_type": "EGNN",
+            "equivariance": true,
+            "radius": 5.0,
+            "max_neighbours": 100000,
+            "num_gaussians": 50,
+            "envelope_exponent": 5,
+            "int_emb_size": 64,
+            "basis_emb_size": 8,
+            "out_emb_size": 128,
+            "num_after_skip": 2,
+            "num_before_skip": 1,
+            "num_radial": 6,
+            "num_spherical": 7,
+            "num_filters": 126,
+            "edge_features": ["length"],
+            "hidden_dim": 50,
+            "num_conv_layers": 3,
+            "output_heads": {
+                "node": {
+                    "num_headlayers": 2,
+                    "dim_headlayers": [200,200],
+                    "type": "mlp"
+                }
+            },
+            "task_weights": [1.0]
+        },
+        "Variables_of_interest": {
+            "input_node_features": [0, 1, 2, 3],
+            "output_names": ["forces"],
+            "output_index": [2],
+            "output_dim": [3],
+            "type": ["node"]
+        },
+        "Training": {
+            "num_epoch": 50,
+            "EarlyStopping": true,
+            "perc_train": 0.9,
+            "loss_function_type": "mae",
+            "batch_size": 32,
+            "continue": 0,
+            "Optimizer": {
+                "type": "AdamW",
+                "learning_rate": 1e-3
+            }
+        }
+    },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
+    }
+}