tfep.nn.graph.compute_edge_distances

tfep.nn.graph.compute_edge_distances(x, edges, normalize_directions=False, inverse_directions=False)[source]

Return distances between nodes across edges.

Parameters:

  • x (torch.Tensor) – Positions of particles with shape (batch_size*n_nodes, 3).

  • edges (torch.Tensor) – Shape (2, batch_size*n_edges). The i-th edge is created from node edges[0][i] to edges[1][i], where edges[0][i] is a node index in the range [0, batch_size*n_nodes].

  • normalize_directions (bool, optional) – If True, the returned directions are normalized by their norms.

  • inverse_directions (bool, optional) – If True, the direction vectors go from the destination node to the source node rather than the opposite.

Returns:

  • distances (torch.Tensor) – Shape (batch_size*n_edges,). distances[i] is the distance between the nodes of the i-th edge.

  • direction (torch.Tensor) – Shape (batch_size*n_edges, 3). direction[i] is the vector connecting the nodes of the i-th edge. The vectors are normalized by their norms if normalize_direction is True.