torch_geometric.nn.models.GNNFF

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

Bases: Module

图神经网络力场（GNNFF）来自 “准确且可扩展的图神经网络力场及直接力架构的分子动力学”论文。 GNNFF 直接从局部原子环境的自动提取特征中预测原子力，这些特征在平移上是不变的，但在原子的坐标上是旋转协变的。

Parameters:

• hidden_node_channels (int) – 隐藏节点嵌入大小。

• hidden_edge_channels (int) – 隐藏边嵌入大小。

• num_layers (int) – 消息传递块的数量。

• 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

Return type:

Tensor

reset_parameters()