Skip to content

Commit

Permalink
adapted hyper for train_force to simplify.
Browse files Browse the repository at this point in the history
  • Loading branch information
@PatReis
PatReis committed Aug 31, 2023
1 parent 6a68efe commit c16e7f8
Show file tree
Hide file tree
Showing 4 changed files with 158 additions and 71 deletions.
9 changes: 7 additions & 2 deletions kgcnn/model/force.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ class EnergyForceModel(ks.models.Model):

def __init__(self,
model_energy=None,
inputs: dict = None,
coordinate_input: Union[int, str] = 1,
output_as_dict: bool = True,
ragged_validate: bool = False,
Expand All @@ -80,6 +81,7 @@ def __init__(self,
Args:
model_energy (dict): Keras model for energy prediction. Can also be a serialization dict.
inputs (dict, list): List or Dict of inputs for the force model.
coordinate_input (str, int): Index or key where to find coordinate tensor in model input.
output_as_dict (bool): Whether to return energy and force as list or as dict. Default is True.
ragged_validate (bool): Whether to validate ragged tensor creation. Default is False.
Expand All @@ -95,7 +97,6 @@ def __init__(self,
super(EnergyForceModel, self).__init__(self, **kwargs)
if model_energy is None:
raise ValueError("Require valid model in `model_energy` for force prediction.")

# Input for model_energy.
self._model_energy = model_energy

Expand All @@ -121,6 +122,9 @@ def __init__(self,
self.nested_model_config = nested_model_config
self.use_batch_jacobian = use_batch_jacobian

# We can try to infer the model inputs from energy model, if not given explicit.
self._inputs_to_force_model = inputs

# Layers.
self.cast_coordinates = ChangeTensorType(input_tensor_type="ragged", output_tensor_type="mask")

Expand Down Expand Up @@ -209,6 +213,7 @@ def get_config(self):
"output_to_tensor": self.output_to_tensor,
"output_squeeze_states": self.output_squeeze_states,
"nested_model_config": self.nested_model_config,
"use_batch_jacobian": self.use_batch_jacobian
"use_batch_jacobian": self.use_batch_jacobian,
"inputs": self._inputs_to_force_model
})
return conf
95 changes: 68 additions & 27 deletions training/hyper/hyper_iso17.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,18 @@
"output_to_tensor": False,
"output_squeeze_states": True,
"coordinate_input": 1,
"inputs": [
{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True}
],
"model_energy": {
"class_name": "make_model",
"module_name": "kgcnn.literature.Schnet",
"config": {
"name": "SchnetEnergy",
"inputs": [
{"shape": [None], "name": "numbers", "dtype": "float32", "ragged": True},
{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True}
],
Expand All @@ -38,9 +43,6 @@
}
},
"training": {
"target_property_names": {
"energy": "total_energy", "force": "atomic_forces", "atomic_number": "numbers",
"coordinates": "positions"},
"train_test_indices": {"train": "train", "test": "test", "split_index": [0]},
"fit": {
"batch_size": 64, "epochs": 1000, "validation_freq": 1, "verbose": 2,
Expand All @@ -62,6 +64,10 @@
"module_name": "kgcnn.data.datasets.ISO17Dataset",
"config": {},
"methods": [
{"rename_property_on_graphs": {"old_property_name": "total_energy", "new_property_name": "energy"}},
{"rename_property_on_graphs": {"old_property_name": "atomic_forces", "new_property_name": "force"}},
{"rename_property_on_graphs": {"old_property_name": "numbers",
"new_property_name": "atomic_number"}},
{"map_list": {"method": "set_range", "max_distance": 5, "max_neighbours": 10000,
"node_coordinates": "positions"}}
]
Expand All @@ -83,13 +89,18 @@
"output_to_tensor": False,
"output_squeeze_states": True,
"coordinate_input": 1,
"inputs": [
{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True}
],
"model_energy": {
"module_name": "kgcnn.literature.PAiNN",
"class_name": "make_model",
"config": {
"name": "PAiNNEnergy",
"inputs": [
{"shape": [None], "name": "numbers", "dtype": "float32", "ragged": True},
{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True}
],
Expand All @@ -105,9 +116,6 @@
}
},
"training": {
"target_property_names": {
"energy": "total_energy", "force": "atomic_forces", "atomic_number": "numbers",
"coordinates": "positions"},
"train_test_indices": {"train": "train", "test": "test", "split_index": [0]},
"fit": {
"batch_size": 32, "epochs": 1000, "validation_freq": 1, "verbose": 2,
Expand Down Expand Up @@ -141,6 +149,10 @@
"config": {
},
"methods": [
{"rename_property_on_graphs": {"old_property_name": "total_energy", "new_property_name": "energy"}},
{"rename_property_on_graphs": {"old_property_name": "atomic_forces", "new_property_name": "force"}},
{"rename_property_on_graphs": {"old_property_name": "numbers",
"new_property_name": "atomic_number"}},
{"map_list": {"method": "set_range", "max_distance": 5, "max_neighbours": 10000,
"node_coordinates": "positions"}}
]
Expand All @@ -162,12 +174,16 @@
"output_to_tensor": False,
"output_squeeze_states": True,
"coordinate_input": 1,
"inputs": [{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True},
{"shape": [None, 2], "name": "angle_indices", "dtype": "int64", "ragged": True}],
"model_energy": {
"class_name": "make_model",
"module_name": "kgcnn.literature.DimeNetPP",
"config": {
"name": "DimeNetPPEnergy",
"inputs": [{"shape": [None], "name": "numbers", "dtype": "float32", "ragged": True},
"inputs": [{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True},
{"shape": [None, 2], "name": "angle_indices", "dtype": "int64", "ragged": True}],
Expand All @@ -189,9 +205,6 @@
}
},
"training": {
"target_property_names": {
"energy": "total_energy", "force": "atomic_forces", "atomic_number": "numbers",
"coordinates": "positions"},
"train_test_indices": {"train": "train", "test": "test", "split_index": [0]},
"fit": {
"batch_size": 10, "epochs": 1000, "validation_freq": 1, "verbose": 2,
Expand Down Expand Up @@ -225,6 +238,10 @@
"config": {
},
"methods": [
{"rename_property_on_graphs": {"old_property_name": "total_energy", "new_property_name": "energy"}},
{"rename_property_on_graphs": {"old_property_name": "atomic_forces", "new_property_name": "force"}},
{"rename_property_on_graphs": {"old_property_name": "numbers",
"new_property_name": "atomic_number"}},
{"map_list": {"method": "set_range", "max_distance": 5, "max_neighbours": 1000,
"node_coordinates": "positions"}},
{"map_list": {"method": "set_angle", "node_coordinates": "positions"}}
Expand All @@ -248,12 +265,15 @@
"output_to_tensor": False,
"output_squeeze_states": True,
"coordinate_input": 1,
"inputs": [{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True}],
"model_energy": {
"class_name": "make_model",
"module_name": "kgcnn.literature.NMPN",
"config": {
"name": "NMPNEnergy",
"inputs": [{"shape": [None], "name": "numbers", "dtype": "float32", "ragged": True},
"inputs": [{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True}],
"input_embedding": {"node": {"input_dim": 95, "output_dim": 64},
Expand All @@ -273,9 +293,6 @@
}
},
"training": {
"target_property_names": {
"energy": "total_energy", "force": "atomic_forces", "atomic_number": "numbers",
"coordinates": "positions"},
"train_test_indices": {"train": "train", "test": "test", "split_index": [0]},
"fit": {
"batch_size": 64, "epochs": 1000, "validation_freq": 1, "verbose": 2,
Expand All @@ -298,6 +315,10 @@
"config": {
},
"methods": [
{"rename_property_on_graphs": {"old_property_name": "total_energy", "new_property_name": "energy"}},
{"rename_property_on_graphs": {"old_property_name": "atomic_forces", "new_property_name": "force"}},
{"rename_property_on_graphs": {"old_property_name": "numbers",
"new_property_name": "atomic_number"}},
{"map_list": {"method": "set_range", "max_distance": 4, "max_neighbours": 1000,
"node_coordinates": "positions"}}
]
Expand All @@ -320,13 +341,19 @@
"output_to_tensor": False,
"output_squeeze_states": True,
"coordinate_input": 1,
"inputs": [
{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True},
{"shape": [3], "name": "graph_inertia", "dtype": "float32", "ragged": False}
],
"model_energy": {
"class_name": "make_model",
"module_name": "kgcnn.literature.Megnet",
"config": {
"name": "MegnetEnergy",
"inputs": [
{"shape": [None], "name": "numbers", "dtype": "float32", "ragged": True},
{"shape": [None], "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": [None, 3], "name": "positions", "dtype": "float32", "ragged": True},
{"shape": [None, 2], "name": "range_indices", "dtype": "int64", "ragged": True},
{"shape": [3], "name": "graph_inertia", "dtype": "float32", "ragged": False}
Expand All @@ -350,9 +377,6 @@
}
},
"training": {
"target_property_names": {
"energy": "total_energy", "force": "atomic_forces", "atomic_number": "numbers",
"coordinates": "positions"},
"train_test_indices": {"train": "train", "test": "test", "split_index": [0]},
"fit": {
"batch_size": 64, "epochs": 1000, "validation_freq": 1, "verbose": 2,
Expand All @@ -375,6 +399,10 @@
"config": {
},
"methods": [
{"rename_property_on_graphs": {"old_property_name": "total_energy", "new_property_name": "energy"}},
{"rename_property_on_graphs": {"old_property_name": "atomic_forces", "new_property_name": "force"}},
{"rename_property_on_graphs": {"old_property_name": "numbers",
"new_property_name": "atomic_number"}},
{"map_list": {"method": "principal_moments_of_inertia", "node_mass": "numbers",
"node_coordinates": "positions"}},
{"map_list": {"method": "set_range", "max_distance": 4, "max_neighbours": 1000,
Expand All @@ -399,12 +427,19 @@
"output_to_tensor": False,
"output_squeeze_states": True,
"coordinate_input": 1,
"inputs": [{"shape": (None,), "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": (None, 3), "name": "positions", "dtype": "float32", "ragged": True},
{"shape": (None, 1), "name": "edge_weights", "dtype": "float32", "ragged": True},
{"shape": (None, 2), "name": "edge_indices", "dtype": "int64", "ragged": True},
{"shape": [None, 2], "name": "angle_indices_1", "dtype": "int64", "ragged": True},
{"shape": [None, 2], "name": "angle_indices_2", "dtype": "int64", "ragged": True},
{"shape": (None, 2), "name": "range_indices", "dtype": "int64", "ragged": True}],
"model_energy": {
"class_name": "make_model",
"module_name": "kgcnn.literature.MXMNet",
"config": {
"name": "MXMNetEnergy",
"inputs": [{"shape": (None,), "name": "numbers", "dtype": "float32", "ragged": True},
"inputs": [{"shape": (None,), "name": "atomic_number", "dtype": "float32", "ragged": True},
{"shape": (None, 3), "name": "positions", "dtype": "float32", "ragged": True},
{"shape": (None, 1), "name": "edge_weights", "dtype": "float32", "ragged": True},
{"shape": (None, 2), "name": "edge_indices", "dtype": "int64", "ragged": True},
Expand Down Expand Up @@ -438,9 +473,6 @@
}
},
"training": {
"target_property_names": {
"energy": "total_energy", "force": "atomic_forces", "atomic_number": "numbers",
"coordinates": "positions"},
"train_test_indices": {"train": "train", "test": "test", "split_index": [0]},
"fit": {
"batch_size": 128, "epochs": 1000, "validation_freq": 1, "verbose": 2,
Expand All @@ -463,6 +495,10 @@
"config": {
},
"methods": [
{"rename_property_on_graphs": {"old_property_name": "total_energy", "new_property_name": "energy"}},
{"rename_property_on_graphs": {"old_property_name": "atomic_forces", "new_property_name": "force"}},
{"rename_property_on_graphs": {"old_property_name": "numbers",
"new_property_name": "atomic_number"}},
# we have to make edges via range cutoff.
{"map_list": {"method": "set_range", "max_distance": 3.0, "max_neighbours": 1000,
"node_coordinates": "positions", "range_indices": "edge_indices",
Expand Down Expand Up @@ -499,6 +535,10 @@
"output_to_tensor": False,
"output_squeeze_states": True,
"coordinate_input": 1,
"inputs": [{"shape": (None, 15), "name": "node_attributes", "dtype": "float32", "ragged": True},
{"shape": (None, 3), "name": "positions", "dtype": "float32", "ragged": True},
{"shape": (None, 2), "name": "range_indices", "dtype": "int64", "ragged": True},
{"shape": (None, 1), "name": "range_attributes", "dtype": "int64", "ragged": True}],
"model_energy": {
"class_name": "make_model",
"module_name": "kgcnn.literature.EGNN",
Expand Down Expand Up @@ -537,9 +577,6 @@
}
},
"training": {
"target_property_names": {
"energy": "total_energy", "force": "atomic_forces", "atomic_number": "numbers",
"coordinates": "positions"},
"train_test_indices": {"train": "train", "test": "test", "split_index": [0]},
"fit": {
"batch_size": 96, "epochs": 1000, "validation_freq": 1, "verbose": 2,
Expand All @@ -562,6 +599,10 @@
"config": {
},
"methods": [
{"rename_property_on_graphs": {"old_property_name": "total_energy", "new_property_name": "energy"}},
{"rename_property_on_graphs": {"old_property_name": "atomic_forces", "new_property_name": "force"}},
{"rename_property_on_graphs": {"old_property_name": "numbers",
"new_property_name": "atomic_number"}},
{"map_list": {"method": "atomic_charge_representation", "node_number": "numbers"}},
{"map_list": {"method": "set_range", "max_distance": 10, "max_neighbours": 10000,
"node_coordinates": "positions"}}
Expand Down
Loading

0 comments on commit c16e7f8

Please sign in to comment.