model

Note

You can customize the PotentialModel to train and predict atom and bond related properties. You need to store the labels on graph edges if you want to do so. This model has multiple attention heads.

class PotentialModel(nn.Module)

A GAT model with multiple gat layers for predicting atomic energies, forces, and stress tensors.

Important

The first value of gat_node_dim_list is the depth of atomic representation.

The first value of energy_readout_node_list, force_readout_node_list, stress_readout_node_list is the input dimension and equals to last value of gat_node_list * num_heads.

The last values of energy_readout_node_list, force_readout_node_list, stress_readout_node_list are 1, 3, and 6, respectively.

__init__(self, gat_node_dim_list, energy_readout_node_list, force_readout_node_list, stress_readout_node_list, head_list=['div'], bias=True, negative_slope=0.2, device='cuda', tail_readout_no_act=[3, 3, 3])
Parameters:

  • gat_node_dim_list (list) – A list of node dimensions of the AGAT Layer.

  • energy_readout_node_list (list) – A list of node dimensions of the energy readout layers.

  • force_readout_node_list (list) – A list of node dimensions of the force readout layers.

  • stress_readout_node_list (list) – A list of node dimensions of the stress readout layers.

  • head_list (list, optional) – A list of attention head names, defaults to [‘div’]

  • bias (TYPE, bool) – Add bias or not to the neural networks., defaults to True

  • negative_slope (float, optional) – This specifies the negative slope of the LeakyReLU (see https://pytorch.org/docs/stable/generated/torch.nn.LeakyReLU.html) activation function., defaults to 0.2

  • device (str, optional) – Device to train the model. Use GPU cards to accerelate training., defaults to ‘cuda’

  • tail_readout_no_act (list, optional) – The tail tail_readout_no_act layers will have no activation functions. The first, second, and third elements are for energy, force, and stress readout layers, respectively., defaults to [3,3,3]

Returns:

An AGAT model

Return type:

agat.model.PotentialModel

mul(self, TorchTensor)

Multiply head.

Parameters:

TorchTensor (torch.tensor) – Input tensor

Returns:

Ouput tensor

Return type:

torch.tensor

div(self, TorchTensor)

Division head.

Parameters:

TorchTensor (torch.tensor) – Input tensor

Returns:

Ouput tensor

Return type:

torch.tensor

div(self, TorchTensor)

Free head.

Parameters:

TorchTensor (torch.tensor) – Input tensor

Returns:

Ouput tensor all ones

Return type:

torch.tensor

get_head_mechanism(self, fn_list, TorchTensor)

Get attention heads

Parameters:

  • fn_list (list) – A list of head mechanisms. For example: [‘mul’, ‘div’, ‘free’]

  • TorchTensor (torch.Tensor) – A PyTorch tensor

Returns:

A new tensor after the transformation.

Return type:

torch.Tensor

forward(self, graph)

The forward function of PotentialModel model.

Parameters:

graph (DGL.Graph) – DGL.Graph

Returns:

  • energy: atomic energy

  • force: atomic force

  • stress: cell stress tensor

Return type:

tuple of torch.tensors