kgcnn.models package¶
Submodules¶
kgcnn.models.casting module¶
-
kgcnn.models.casting.
template_cast_dict_input
(model_inputs: dict, input_tensor_type: str, cast_disjoint_kwargs: dict, mask_assignment: Optional[dict] = None, index_assignment: Optional[dict] = None, return_sub_id: bool = True, rename_mask_to_id: Optional[dict] = None)[source]¶ Template to cast a dictionary of model inputs to a dict of disjoint tensors. The number of model inputs can be variable. The system is rather flexible and explained by
template_cast_dict_input_docs
.The ID information from the mask in tensor form is appended to the returned dictionary. Ragged Tensor do not require a mask but use the mask information to generate joint ID tensors.
- Parameters
model_inputs (dict) – Dictionary of Keras inputs.
input_tensor_type (str) – Input tensor type. Either “padded”, “ragged” or “disjoint”.
cast_disjoint_kwargs (dict) – Kwargs for casting layers.
mask_assignment (dict) – Dictionary of mask name for each input that requires a mask and is cast to a disjoint tensor representation.
index_assignment (dict) – Dictionary of assigning indices to the name of their target tensors to which the indices refer to.
return_sub_id (bool) – Whether the returned dict contains the sub-graph ID tensors.
rename_mask_to_id (dict) – Mapping of mask names to ID names.
- Returns
Model input tensors in disjoint representation.
- Return type
-
kgcnn.models.casting.
template_cast_list_input
(model_inputs: list, input_tensor_type: str, cast_disjoint_kwargs: dict, mask_assignment: Optional[list] = None, index_assignment: Optional[list] = None, return_sub_id: bool = True)[source]¶ Template to cast a list of model inputs to a list of disjoint tensors. The number of model inputs can be variable. The system is preset and explained by
template_cast_list_input_docs
.The ID information from the mask in tensor form is appended to the returned list. Ragged Tensor do not require a mask but use the mask information to generate joint ID tensors.
- Parameters
model_inputs (list) – List of Keras inputs.
input_tensor_type (str) – Input tensor type. Either “padded”, “ragged” or “disjoint”.
cast_disjoint_kwargs (dict) – Kwargs for casting layers.
mask_assignment (list) – List that assigns Tensors to their mask. Inputs that do not require a mask must be marked by Nones. Different inputs can use the same mask.
index_assignment (list) – List that assigns index Tensors to their target to which the index refer to. Inputs that are not indices must be marked by Nones.
return_sub_id (bool) – Whether the returned list contains the sub-graph ID tensors.
- Returns
List of Keras Tensors for disjoint model.
- Return type
-
kgcnn.models.casting.
template_cast_output
(model_outputs, output_embedding, output_tensor_type, input_tensor_type, cast_disjoint_kwargs)[source]¶ Template to cast graph, node or edge output to the desired tensor representation.
- Parameters
model_outputs (list) – List of output and additional ID tensors. The list must always be [model_output, batch_id_node, batch_id_edge, node_id, edge_id, count_nodes, count_edges] but can have None in place of ID tensors if not required.
output_embedding (str) – Embedding of the graph output, either “graph”, “node” or “edge”.
output_tensor_type (str) – The tensor representation of model output.
input_tensor_type (str) – The tensor representation of model input.
cast_disjoint_kwargs (dict) – Kwargs for casting layers.
- Returns
Keras output tensor.
- Return type
Tensor
kgcnn.models.force module¶
-
class
kgcnn.models.force.
EnergyForceModel
(*args, **kwargs)[source]¶ Bases:
keras.src.models.model.Model
Outer model to wrap a normal invariant GNN to predict forces from energy predictions via partial derivatives. The Force \(\vec{F_i}\) on Atom \(i\) is given by
\[\vec{F_i} = - \vec{\nabla}_i E_{\text{total}}\]Note that the model simply returns the tensor type of the coordinate input for forces. No casting is done by this class. This means that the model returns a ragged, disjoint or padded tensor depending on the tensor type of the coordinates.
-
__init__
(inputs: Optional[Union[list, dict]] = None, model_energy=None, coordinate_input: Union[int, str] = 1, output_as_dict: bool = True, ragged_validate: bool = False, output_to_tensor: bool = True, output_squeeze_states: bool = False, nested_model_config: bool = True, is_physical_force: bool = True, use_batch_jacobian: Optional[bool] = None, name: Optional[str] = None, outputs: Optional[Union[list, dict]] = None)[source]¶ Initialize Force model with an energy model.
- Parameters
inputs (list) – List of inputs as dictionary kwargs of keras input layers.
model_energy (ks.models.Model, dict) – Keras model os deserialization dictionary for a keras model.
coordinate_input (int) – Position of the coordinate input.
output_as_dict (bool, tuple) – Names for the output if a dictionary should be returned. Or simply a bool which will use the names “energy” and “force”.
ragged_validate (bool) – Whether to validate ragged or jagged tensors.
output_to_tensor – Deprecated.
output_squeeze_states (bool) – Whether to squeeze state/energy dimension for forces in case of a single energy value.
nested_model_config (bool) – Whether there is a config for the energy model.
is_physical_force (bool) – Whether to return the physical force, e.g. the negative gradient of the energy.
use_batch_jacobian – Deprecated.
name (str) – Name of the model.
outputs – List of outputs as dictionary kwargs similar to inputs. Not used by the model but can be passed for external use.
-
kgcnn.models.multi module¶
kgcnn.models.serial module¶
kgcnn.models.utils module¶
-
kgcnn.models.utils.
change_attributes_in_all_layers
(model, attributes_to_change=None, layer_type=None)[source]¶ Utility/helper function to change the attributes from a dictionary in all layers of a model of a certain type.
Warning
This function can change attributes but can cause problems for built models. Also take care which attributes you are changing, especially if they include weights. Always check model behaviour after applying this function.
- Parameters
model (tf.keras.models.Model) – Model to modify.
attributes_to_change (dict) – Dictionary of attributes to change in all layers of a specific type.
layer_type – Class type of the layer to change. Default is None.
- Returns
Model which has layers with changed attributes.
- Return type
tf.keras.models.Model
-
kgcnn.models.utils.
get_model_class
(module_name: str, class_name: str)[source]¶ Helper function to get model class by string identifier.
-
kgcnn.models.utils.
update_model_kwargs
(model_default, update_recursive=inf, deprecated: Optional[list] = None)[source]¶ Decorating function for update_model_kwargs_logic() .
-
kgcnn.models.utils.
update_model_kwargs_logic
(default_kwargs: Optional[dict] = None, user_kwargs: Optional[dict] = None, update_recursive: Union[int, float] = inf)[source]¶ Make model kwargs dictionary with updated default values. This is essentially a nested version of update() for dicts. This is supposed to be more convenient if the values of kwargs are again layer kwargs to be unpacked, and do not need to be fully known to update them.
- Parameters
- Returns
New dict and update with first default and then user args.
- Return type