kgcnn.literature.DimeNetPP package¶

Module contents¶

kgcnn.literature.DimeNetPP.make_crystal_model(inputs: list = None, input_tensor_type: str = None, cast_disjoint_kwargs: dict = None, input_embedding: dict = None, input_node_embedding: dict = None, emb_size: int = None, out_emb_size: int = None, int_emb_size: int = None, basis_emb_size: int = None, num_blocks: int = None, num_spherical: int = None, num_radial: int = None, cutoff: float = None, envelope_exponent: int = None, num_before_skip: int = None, num_after_skip: int = None, num_dense_output: int = None, num_targets: int = None, activation: str = None, extensive: bool = None, output_init: str = None, verbose: int = None, name: str = None, output_embedding: str = None, output_tensor_type: str = None, use_output_mlp: bool = None, output_mlp: dict = None, output_scaling: dict = None)[source]¶

Make DimeNetPP graph network via functional API. Default parameters can be found in kgcnn.literature.DimeNetPP.model_crystal_default.

Note

DimeNetPP does require a large amount of memory for this implementation, which increase quickly with the number of connections in a batch. Use ragged input or dataloader if possible.

Model inputs: Model uses the list template of inputs and standard output template. The supported inputs are [nodes, coordinates, edge_indices, angle_indices, image_translation, lattice, ...] with ‘…’ indicating mask or ID tensors following the template below. Note that you must supply angle indices as index pairs that refer to two edges.

Template of listed graph input tensors, which should be compatible to previous kgcnn versions and defines the order as follows: [nodes, edges, angles, edge_indices, angle_indices, graph_state, ...] . Where ‘…’ denotes further mask or ID tensors, which is required for certain input types (see below). Depending on the model, some inputs may not be used (see model description for information on supported inputs). For example if the model does not support angles and no graph attribute input, the input becomes: [nodes, edges, edge_indices, ...] . In case of crystal graphs lattice and translation information has to be added. This will give a possible input of [nodes, edges, angles, edge_indices, angle_indices, graph_state, image_translation, lattice,...] . Note that in place of nodes or edges also more than one tensor can be provided, depending on the model, for example [nodes_1, nodes_2, edges_1, edges_2, edge_indices, ...] .

However, for future models we intend to used named inputs rather than a list that is sensible to ordering. Whether to use mask or length tensor for padded as well as further parameter of casting has to be set with (dict) cast_disjoint_kwargs .

Padded or Masked Inputs:

list: [nodes, edges, angles, edge_indices, angle_indices, graph_state, image_translation, lattice, node_mask/node_count, edge_mask/edge_count, angle_mask/angle_count]

nodes (Tensor): Node attributes of shape (batch, N, F) or (batch, N) using an embedding layer.

edges (Tensor): Edge attributes of shape (batch, M, F) or (batch, M) using an embedding layer.

angles (Tensor): Angle attributes of shape (batch, M, F) or (batch, K) using an embedding layer.

edge_indices (Tensor): Index list for edges of shape (batch, M, 2) referring to nodes.

angle_indices (Tensor): Index list for angles of shape (batch, K, 2) referring to edges.

graph_state (Tensor): Graph attributes of shape (batch, F) .

image_translation (Tensor): Indices of the periodic image the sending node is located in. Shape is (batch, M, 3) .

lattice (Tensor): Lattice matrix of the periodic structure of shape (batch, 3, 3) .

node_mask (Tensor): Mask for padded nodes of shape (batch, N) .

edge_mask (Tensor): Mask for padded edges of shape (batch, M) .

angle_mask (Tensor): Mask for padded angles of shape (batch, K) .

node_count (Tensor): Total number of nodes if padding is used of shape (batch, ) .

edge_count (Tensor): Total number of edges if padding is used of shape (batch, ) .

angle_count (Tensor): Total number of angle if padding is used of shape (batch, ) .

Ragged or Jagged Inputs:

list: [nodes, edges, angles, edge_indices, angle_indices, graph_state, image_translation, lattice]

nodes (RaggedTensor): Node attributes of shape (batch, None, F) or (batch, None) using an embedding layer.

edges (RaggedTensor): Edge attributes of shape (batch, None, F) or (batch, None) using an embedding layer.

angles (RaggedTensor): Angle attributes of shape (batch, None, F) or (batch, None) using an embedding layer.

edge_indices (RaggedTensor): Index list for edges of shape (batch, None, 2) referring to nodes.

angle_indices (RaggedTensor): Index list for angles of shape (batch, None, 2) referring to edges.

graph_state (Tensor): Graph attributes of shape (batch, F) .

image_translation (RaggedTensor): Indices of the periodic image the sending node is located in. Shape is (batch, None, 3) .

lattice (Tensor): Lattice matrix of the periodic structure of shape (batch, 3, 3) .

Disjoint Input:

list: [nodes, edges, angles, edge_indices, angle_indices, graph_state, image_translation, lattice, graph_id_node, graph_id_edge, graph_id_angle, nodes_id, edges_id, angle_id, nodes_count, edges_count, angles_count]

nodes (Tensor): Node attributes of shape ([N], F) or ([N], ) using an embedding layer.

edges (Tensor): Edge attributes of shape ([M], F) or ([M], ) using an embedding layer.

angles (Tensor): Angle attributes of shape ([K], F) or ([K], ) using an embedding layer.

edge_indices (Tensor): Index list for edges of shape (2, [M]) referring to nodes.

angle_indices (Tensor): Index list for angles of shape (2, [K]) referring to edges.

graph_state (Tensor): Graph attributes of shape (batch, F) .

image_translation (Tensor): Indices of the periodic image the sending node is located in. Shape is ([M], 3) .

lattice (Tensor): Lattice matrix of the periodic structure of shape (batch, 3, 3) .

graph_id_node (Tensor): ID tensor of batch assignment in disjoint graph of shape ([N], ) .

graph_id_edge (Tensor): ID tensor of batch assignment in disjoint graph of shape ([M], ) .

graph_id_angle (Tensor): ID tensor of batch assignment in disjoint graph of shape ([K], ) .

nodes_id (Tensor): The ID-tensor to assign each node to its respective graph of shape ([N], ) .

edges_id (Tensor): The ID-tensor to assign each edge to its respective graph of shape ([M], ) .

angle_id (Tensor): The ID-tensor to assign each edge to its respective graph of shape ([K], ) .

nodes_count (Tensor): Tensor of number of nodes for each graph of shape (batch, ) .

edges_count (Tensor): Tensor of number of edges for each graph of shape (batch, ) .

angles_count (Tensor): Tensor of number of angles for each graph of shape (batch, ) .

Model outputs: The standard output template:

The standard model output template returns a single tensor of either “graph”, “node”, or “edge” embeddings specified by output_embedding within the model. The return tensor type is determined by output_tensor_type . Options are:

graph:

Tensor: Graph labels of shape (batch, F) .

nodes:

Tensor: Node labels for the graph of either type:

ragged (RaggedTensor): Single tensor of shape (batch, None, F) .

padded (Tensor): Padded tensor of shape (batch, N, F) .

disjoint (Tensor): Disjoint representation of shape ([N], F) .

edges:

Tensor: Edge labels for the graph of either type:

ragged (RaggedTensor): Single tensor of shape (batch, None, F) .

padded (Tensor): Padded tensor of shape (batch, M, F)

disjoint (Tensor): Disjoint representation of shape ([M], F) .

Parameters

inputs (list) – List of dictionaries unpacked in tf.keras.layers.Input. Order must match model definition.
input_tensor_type (str) – Input type of graph tensor. Default is “padded”.
cast_disjoint_kwargs (dict) – Dictionary of arguments for casting layer.
input_embedding (dict) – Deprecated in favour of input_node_embedding etc.
input_node_embedding (dict) – Dictionary of embedding arguments for nodes unpacked in Embedding layers.
emb_size (int) – Overall embedding size used for the messages.
out_emb_size (int) – Embedding size for output of DimNetOutputBlock.
int_emb_size (int) – Embedding size used for interaction triplets.
basis_emb_size (int) – Embedding size used inside the basis transformation.
num_blocks (int) – Number of graph embedding blocks or depth of the network.
num_spherical (int) – Number of spherical components in SphericalBasisLayer.
num_radial (int) – Number of radial components in basis layer.
cutoff (float) – Distance cutoff for basis layer.
envelope_exponent (int) – Exponent in envelope function for basis layer.
num_before_skip (int) – Number of residual layers in interaction block before skip connection
num_after_skip (int) – Number of residual layers in interaction block after skip connection
num_dense_output (int) – Number of dense units in output DimNetOutputBlock.
num_targets (int) – Number of targets or output embedding dimension of the model.
activation (str, dict) – Activation to use.
extensive (bool) – Graph output for extensive target to apply sum for pooling or mean otherwise.
output_init (str, dict) – Output initializer for kernel.
verbose (int) – Level of verbosity.
name (str) – Name of the model.
output_embedding (str) – Main embedding task for graph network. Either “node”, “edge” or “graph”.
use_output_mlp (bool) – Whether to use the final output MLP. Possibility to skip final MLP.
output_mlp (dict) – Dictionary of layer arguments unpacked in the final classification MLP layer block. Defines number of model outputs and activation. Note that DimeNetPP originally defines the output dimension via num_targets. But this can be set to out_emb_size and the output_mlp be used for more specific control.
output_scaling (dict) – Dictionary of layer arguments unpacked in scaling layers. Default is None.
output_tensor_type (str) – Output type of graph tensors such as nodes or edges. Default is “padded”.

Returns

keras.models.Model

kgcnn.literature.DimeNetPP.make_model(inputs: list = None, input_tensor_type: str = None, cast_disjoint_kwargs: dict = None, input_embedding: dict = None, input_node_embedding: dict = None, emb_size: int = None, out_emb_size: int = None, int_emb_size: int = None, basis_emb_size: int = None, num_blocks: int = None, num_spherical: int = None, num_radial: int = None, cutoff: float = None, envelope_exponent: int = None, num_before_skip: int = None, num_after_skip: int = None, num_dense_output: int = None, num_targets: int = None, activation: str = None, extensive: bool = None, output_init: str = None, verbose: int = None, name: str = None, output_embedding: str = None, output_tensor_type: str = None, use_output_mlp: bool = None, output_mlp: dict = None, output_scaling: dict = None)[source]¶

Make DimeNetPP graph network via functional API. Default parameters can be found in kgcnn.literature.DimeNetPP.model_default.

Model inputs: Model uses the list template of inputs and standard output template. The supported inputs are [nodes, coordinates, edge_indices, angle_indices...] with ‘…’ indicating mask or ID tensors following the template below. Note that you must supply angle indices as index pairs that refer to two edges.

Template of listed graph input tensors, which should be compatible to previous kgcnn versions and defines the order as follows: [nodes, edges, angles, edge_indices, angle_indices, graph_state, ...] . Where ‘…’ denotes further mask or ID tensors, which is required for certain input types (see below). Depending on the model, some inputs may not be used (see model description for information on supported inputs). For example if the model does not support angles and no graph attribute input, the input becomes: [nodes, edges, edge_indices, ...] . In case of crystal graphs lattice and translation information has to be added. This will give a possible input of [nodes, edges, angles, edge_indices, angle_indices, graph_state, image_translation, lattice,...] . Note that in place of nodes or edges also more than one tensor can be provided, depending on the model, for example [nodes_1, nodes_2, edges_1, edges_2, edge_indices, ...] .