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)

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

dict

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)

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

list

kgcnn.models.casting.template_cast_output(model_outputs, output_embedding, output_tensor_type, input_tensor_type, cast_disjoint_kwargs)

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)

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)

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.

build(input_shape)

call(inputs, training=False, **kwargs)

get_config()

Get config.

kgcnn.models.multi module

kgcnn.models.multi.merge_models(model_list: list, merge_type: str = 'concat', output_mlp: Optional[dict] = None)

Merge a list of models by combining their output.

Parameters

• model_list (list) – List of graph models.

• merge_type (str) – How to merge the output.

• output_mlp (dict) – Kwargs of the final MLP after the models’ output.

Returns

ks.models.Model

kgcnn.models.serial module

kgcnn.models.serial.deserialize(obj_dict: dict) → keras.src.models.model.Model

kgcnn.models.utils module

kgcnn.models.utils.change_attributes_in_all_layers(model, attributes_to_change=None, layer_type=None)

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)

Helper function to get model class by string identifier.

Parameters

• module_name (str) – Name of the module of the model.

• class_name (str) – Name of the model class.

Returns

tf.keras.models.Model

kgcnn.models.utils.update_model_kwargs(model_default, update_recursive=inf, deprecated: Optional[list] = None)

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)

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

• default_kwargs (dict) – Dictionary of default values. Default is None.

• user_kwargs (dict) – Dictionary of args to update. Default is None.

• update_recursive (int) – Max depth to update mappings like dict. Default is inf.

Returns

New dict and update with first default and then user args.

Return type

dict

Module contents