class GNNFF(hidden_node_channels: int, hidden_edge_channels: int, num_layers: int, cutoff: float = 5.0, max_num_neighbors: int = 32)[source]

Bases: Module

The Graph Neural Network Force Field (GNNFF) from the “Accurate and scalable graph neural network force field and molecular dynamics with direct force architecture” paper. GNNFF directly predicts atomic forces from automatically extracted features of the local atomic environment that are translationally-invariant, but rotationally-covariant to the coordinate of the atoms.

  • hidden_node_channels (int) – Hidden node embedding size.

  • hidden_edge_channels (int) – Hidden edge embedding size.

  • num_layers (int) – Number of message passing blocks.

  • cutoff (float, optional) – Cutoff distance for interatomic interactions. (default: 5.0)

  • max_num_neighbors (int, optional) – The maximum number of neighbors to collect for each node within the cutoff distance. (default: 32)

forward(z: Tensor, pos: Tensor, batch: Optional[Tensor] = None) Tensor[source]