Merge pull request #105 from allaffa/dimension_graph_outputs

Dimension graph outputs & CFG loader & optional plotting

examples/eam/NiNb_EAM_bulk.json

@@ -0,0 +1,66 @@
+{
+    "Verbosity": {
+        "level": 2
+    },
+    "Dataset": {
+        "name": "NiNb",
+        "path": {"total": "BULK_DATA"},
+        "format": "CFG",
+        "compositional_stratified_splitting": true,
+	"rotational_invariance": true,
+        "node_features": {
+            "name": ["num_of_protons", "atomic_masses", "atomic_energy"],
+            "dim": [1,1,1],
+            "column_index": [0,1,2]
+        },
+        "graph_features":{
+            "name": ["bulk_modulus"],
+            "dim": [1],
+            "column_index": [2]
+        }
+    },
+    "NeuralNetwork": {
+        "Architecture": {
+            "model_type": "PNA",
+            "radius": 3,
+            "max_neighbours": 100000,
+	    "edge_features": ["lengths"],
+	    "periodic_boundary_conditions": true,
+            "hidden_dim": 50,
+            "num_conv_layers": 10,
+            "output_heads": {
+                "graph":{
+                    "num_sharedlayers": 2,
+                    "dim_sharedlayers": 5,
+                    "num_headlayers": 2,
+                    "dim_headlayers": [50,25]
+                },
+                "node": {
+                    "num_headlayers": 2,
+                    "dim_headlayers": [50,25],
+                    "type": "mlp"
+                }
+            },
+            "task_weights": [1.0]
+        },
+        "Variables_of_interest": {
+            "input_node_features": [0],
+            "output_names": ["bulk_modulus"],
+            "output_index": [0],
+            "type": ["graph"],
+            "denormalize_output": true
+        },
+        "Training": {
+            "num_epoch": 2,
+            "perc_train": 0.7,
+            "learning_rate": 1e-3,
+            "batch_size": 64,
+            "continue": 0
+        }
+    },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
+    }
+}

examples/eam/NiNb_EAM_bulk_multitask.json

@@ -0,0 +1,67 @@
+{
+    "Verbosity": {
+        "level": 2
+    },
+    "Dataset": {
+        "name": "NiNb",
+        "path": {"total": "FCC_Bulk_small"},
+        "format": "CFG",
+        "compositional_stratified_splitting": true,
+	"rotational_invariance": true,
+        "node_features": {
+            "name": ["num_of_protons", "atomic_mass", "atomic_energy", "atomic_force"],
+            "dim": [1,1,1,3],
+            "column_index": [0,1,2,3]
+        },
+        "graph_features":{
+            "name": ["bulk_modulus"],
+            "dim": [1],
+            "column_index": [2]
+        }
+    },
+    "NeuralNetwork": {
+        "Architecture": {
+            "model_type": "PNA",
+            "radius": 3,
+            "max_neighbours": 100000,
+	    "edge_features": ["lengths"],
+	    "periodic_boundary_conditions": true,
+            "hidden_dim": 50,
+            "num_conv_layers": 6,
+            "output_heads": {
+                "graph":{
+                    "num_sharedlayers": 2,
+                    "dim_sharedlayers": 5,
+                    "num_headlayers": 2,
+                    "dim_headlayers": [50,25]
+                },
+                "node": {
+                    "num_headlayers": 2,
+                    "dim_headlayers": [50,25],
+                    "type": "mlp"
+                }
+            },
+            "task_weights": [1.0, 1.0, 1.0]
+        },
+        "Variables_of_interest": {
+            "input_node_features": [0],
+            "output_names": ["bulk_modulus","atomic_energy","atomic_forces"],
+            "output_index": [0, 2, 3],
+            "type": ["graph","node","node"],
+            "denormalize_output": true
+        },
+        "Training": {
+            "num_epoch": 2,
+            "perc_train": 0.7,
+            "learning_rate": 1e-3,
+            "batch_size": 64,
+            "continue": 0,
+            "startfrom": "existing_model"
+        }
+    },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
+    }
+}

examples/eam/NiNb_EAM_energy.json

@@ -0,0 +1,66 @@
+{
+    "Verbosity": {
+        "level": 2
+    },
+    "Dataset": {
+        "name": "NiNb",
+        "path": {"total": "FCC_Ni_Nb"},
+        "format": "CFG",
+        "compositional_stratified_splitting": true,
+	"rotational_invariance": true,
+        "node_features": {
+            "name": ["num_of_protons", "atomic_masses", "atomic_energy"],
+            "dim": [1,1,1],
+            "column_index": [0,1,2]
+        },
+        "graph_features":{
+            "name": [],
+            "dim": [],
+            "column_index": []
+        }
+    },
+    "NeuralNetwork": {
+        "Architecture": {
+            "model_type": "PNA",
+            "radius": 3,
+            "max_neighbours": 100000,
+	    "edge_features": ["lengths"],
+	    "periodic_boundary_conditions": true,
+            "hidden_dim": 50,
+            "num_conv_layers": 10,
+            "output_heads": {
+                "graph":{
+                    "num_sharedlayers": 2,
+                    "dim_sharedlayers": 5,
+                    "num_headlayers": 2,
+                    "dim_headlayers": [50,25]
+                },
+                "node": {
+                    "num_headlayers": 2,
+                    "dim_headlayers": [50,25],
+                    "type": "mlp"
+                }
+            },
+            "task_weights": [1.0]
+        },
+        "Variables_of_interest": {
+            "input_node_features": [0],
+            "output_names": ["atomic_energy"],
+            "output_index": [2],
+            "type": ["node"],
+            "denormalize_output": true
+        },
+        "Training": {
+            "num_epoch": 200,
+            "perc_train": 0.7,
+            "learning_rate": 1e-3,
+            "batch_size": 64,
+            "continue": 0
+        }
+    },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
+    }
+}

examples/eam/NiNb_EAM_multitask.json

@@ -0,0 +1,67 @@
+{
+    "Verbosity": {
+        "level": 2
+    },
+    "Dataset": {
+        "name": "NiNb",
+        "path": {"total": "FCC_Ni_Nb"},
+        "format": "CFG",
+        "compositional_stratified_splitting": true,
+	"rotational_invariance": true,
+        "node_features": {
+            "name": ["num_of_protons", "atomic_mass", "atomic_energy", "atomic_force"],
+            "dim": [1,1,1,3],
+            "column_index": [0,1,2,3]
+        },
+        "graph_features":{
+            "name": [],
+            "dim": [],
+            "column_index": []
+        }
+    },
+    "NeuralNetwork": {
+        "Architecture": {
+            "model_type": "PNA",
+            "radius": 3,
+            "max_neighbours": 100000,
+	    "edge_features": ["lengths"],
+	    "periodic_boundary_conditions": true,
+            "hidden_dim": 50,
+            "num_conv_layers": 6,
+            "output_heads": {
+                "graph":{
+                    "num_sharedlayers": 2,
+                    "dim_sharedlayers": 5,
+                    "num_headlayers": 2,
+                    "dim_headlayers": [50,25]
+                },
+                "node": {
+                    "num_headlayers": 2,
+                    "dim_headlayers": [50,25],
+                    "type": "mlp"
+                }
+            },
+            "task_weights": [1.0, 1.0]
+        },
+        "Variables_of_interest": {
+            "input_node_features": [0],
+            "output_names": ["atomic_energy","atomic_forces"],
+            "output_index": [2, 3],
+            "type": ["node","node"],
+            "denormalize_output": true
+        },
+        "Training": {
+            "num_epoch": 20,
+            "perc_train": 0.7,
+            "learning_rate": 1e-3,
+            "batch_size": 64,
+            "continue": 0,
+            "startfrom": "existing_model"
+        }
+    },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
+    }
+}

examples/eam/eam.py

@@ -0,0 +1,5 @@
+import os
+import hydragnn
+
+filepath = os.path.join(os.path.dirname(__file__), "NiNb_EAM_bulk_multitask.json")
+hydragnn.run_training(filepath)

examples/ising_model/ising_model.json

@@ -54,5 +54,10 @@
             "batch_size": 64,
             "continue": 0
         }
+    },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
     }
 }

examples/lsms/lsms.json

@@ -4,9 +4,10 @@
     },
     "Dataset": {
         "name": "FePt_32atoms",
-        "path": {"total": "FePt"},
+        "path": {"total": "FePt_enthalpy"},
         "format": "LSMS",
         "compositional_stratified_splitting": true,
+        "rotational_invariance": true,
         "node_features": {
             "name": ["num_of_protons","charge_density", "magnetic_moment"],
             "dim": [1,1,1],
@@ -56,5 +57,10 @@
             "continue": 0,
             "startfrom": "existing_model"
         }
+    },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
     }
 }

examples/md17/md17.json

@@ -39,5 +39,10 @@
             "continue": 0,
             "startfrom": "existing_model"
         }
+     },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
     }
 }

examples/md17/md17.py

@@ -17,9 +17,13 @@ def md17_pre_transform(data):
     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]
     hydragnn.preprocess.update_predicted_values(
         var_config["type"],
         var_config["output_index"],
+        graph_features_dim,
+        node_feature_dim,
         data,
     )
     data = compute_edges(data)

examples/qm9/qm9.json

@@ -39,5 +39,10 @@
             "continue": 0,
             "startfrom": "existing_model"
         }
+    },
+    "Visualization": {
+        "plot_init_solution": true,
+        "plot_hist_solution": false,
+        "create_plots": true
     }
 }

examples/qm9/qm9.py

@@ -17,9 +17,13 @@ def qm9_pre_transform(data):
     data.x = data.z.float().view(-1, 1)
     # Only predict free energy (index 10 of 19 properties) for this run.
     data.y = data.y[:, 10] / len(data.x)
+    graph_features_dim = [1]
+    node_feature_dim = [1]
     hydragnn.preprocess.update_predicted_values(
         var_config["type"],
         var_config["output_index"],
+        graph_features_dim,
+        node_feature_dim,
         data,
     )
     device = hydragnn.utils.get_device()

hydragnn/postprocess/postprocess.py

@@ -15,12 +15,12 @@ def output_denormalize(y_minmax, true_values, predicted_values):
         ymin = y_minmax[ihead][0]
         ymax = y_minmax[ihead][1]
         for isample in range(len(predicted_values[ihead])):
-            for iatom in range(len(predicted_values[ihead][isample])):
-                predicted_values[ihead][isample][iatom] = (
-                    predicted_values[ihead][isample][iatom] * (ymax - ymin) + ymin
+            for ifeat in range(len(predicted_values[ihead][isample])):
+                predicted_values[ihead][isample][ifeat] = (
+                    predicted_values[ihead][isample][ifeat] * (ymax - ymin) + ymin
                 )
-                true_values[ihead][isample][iatom] = (
-                    true_values[ihead][isample][iatom] * (ymax - ymin) + ymin
+                true_values[ihead][isample][ifeat] = (
+                    true_values[ihead][isample][ifeat] * (ymax - ymin) + ymin
                 )
 
     return true_values, predicted_values

hydragnn/preprocess/__init__.py

@@ -2,6 +2,8 @@
 
 from .utils import (
     check_if_graph_size_variable,
+    get_radius_graph,
+    get_radius_graph_pbc,
     get_radius_graph_config,
     get_radius_graph_pbc_config,
     RadiusGraphPBC,