kgcnn.literature.Megnet package

Module contents

kgcnn.literature.Megnet.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, input_graph_embedding: dict = None, expand_distance: bool = None, make_distance: bool = None, gauss_args: dict = None, meg_block_args: dict = None, set2set_args: dict = None, node_ff_args: dict = None, edge_ff_args: dict = None, state_ff_args: dict = None, use_set2set: bool = None, nblocks: int = None, has_ff: bool = None, dropout: float = None, name: str = None, verbose: int = None, output_embedding: str = None, output_mlp: dict = None, output_tensor_type: str = None, output_scaling: dict = None)

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

Model inputs: Model uses the list template of inputs and standard output template. The supported inputs are [nodes, coordinates, edge_indices, graph_state, image_translation, lattice, ...] with ‘…’ indicating mask or ID tensors following the template below.

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 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.

• input_graph_embedding (dict) – Dictionary of embedding arguments for graph unpacked in Embedding layers.

• make_distance (bool) – Whether input is distance or coordinates at in place of edges.

• expand_distance (bool) – If the edge input are actual edges or node coordinates instead that are expanded to form edges with a gauss distance basis given edge indices. Expansion uses gauss_args.

• gauss_args (dict) – Dictionary of layer arguments unpacked in GaussBasisLayer layer.

• meg_block_args (dict) – Dictionary of layer arguments unpacked in MEGnetBlock layer.

• set2set_args (dict) – Dictionary of layer arguments unpacked in :obj:PoolingSet2SetEncoder layer.

• node_ff_args (dict) – Dictionary of layer arguments unpacked in MLP feed-forward layer.

• edge_ff_args (dict) – Dictionary of layer arguments unpacked in MLP feed-forward layer.

• state_ff_args (dict) – Dictionary of layer arguments unpacked in MLP feed-forward layer.

• use_set2set (bool) – Whether to use PoolingSet2SetEncoder layer.

• nblocks (int) – Number of graph embedding blocks or depth of the network.

• has_ff (bool) – Use feed-forward MLP in each block.

• dropout (int) – Dropout to use. Default is None.

• name (str) – Name of the model.

• verbose (int) – Verbosity level of print.

• output_embedding (str) – Main embedding task for graph network. Either “node”, “edge” or “graph”.

• output_mlp (dict) – Dictionary of layer arguments unpacked in the final classification MLP layer block. Defines number of model outputs and activation.

• 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.Megnet.make_model(inputs: list = None, input_tensor_type: str = None, cast_disjoint_kwargs: dict = None, input_embedding: dict = None, input_node_embedding: dict = None, input_graph_embedding: dict = None, expand_distance: bool = None, make_distance: bool = None, gauss_args: dict = None, meg_block_args: dict = None, set2set_args: dict = None, node_ff_args: dict = None, edge_ff_args: dict = None, state_ff_args: dict = None, use_set2set: bool = None, nblocks: int = None, has_ff: bool = None, dropout: float = None, name: str = None, verbose: int = None, output_embedding: str = None, output_mlp: dict = None, output_tensor_type: str = None, output_scaling: dict = None)

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

Model inputs: Model uses the list template of inputs and standard output template. The supported inputs are [nodes, coordinates, edge_indices, graph_state, ...] with make_distance and with ‘…’ indicating mask or ID tensors following the template below. Note that you could also supply edge features with make_distance to False, which would make the input [nodes, edges, edge_indices, graph_state...] .

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 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.

• input_graph_embedding (dict) – Dictionary of embedding arguments for graph unpacked in Embedding layers.

• make_distance (bool) – Whether input is distance or coordinates at in place of edges.

• expand_distance (bool) – If the edge input are actual edges or node coordinates instead that are expanded to form edges with a gauss distance basis given edge indices. Expansion uses gauss_args.

• gauss_args (dict) – Dictionary of layer arguments unpacked in GaussBasisLayer layer.

• meg_block_args (dict) – Dictionary of layer arguments unpacked in MEGnetBlock layer.

• set2set_args (dict) – Dictionary of layer arguments unpacked in :obj:PoolingSet2SetEncoder layer.

• node_ff_args (dict) – Dictionary of layer arguments unpacked in MLP feed-forward layer.

• edge_ff_args (dict) – Dictionary of layer arguments unpacked in MLP feed-forward layer.

• state_ff_args (dict) – Dictionary of layer arguments unpacked in MLP feed-forward layer.

• use_set2set (bool) – Whether to use PoolingSet2SetEncoder layer.

• nblocks (int) – Number of graph embedding blocks or depth of the network.

• has_ff (bool) – Use feed-forward MLP in each block.

• dropout (int) – Dropout to use. Default is None.

• name (str) – Name of the model.

• verbose (int) – Verbosity level of print.

• output_embedding (str) – Main embedding task for graph network. Either “node”, “edge” or “graph”.

• output_mlp (dict) – Dictionary of layer arguments unpacked in the final classification MLP layer block. Defines number of model outputs and activation.

• 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