core.graph.compute

core.graph.compute#

Copyright (c) Meta Platforms, Inc. and affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Functions#

get_pbc_distances(pos, edge_index, cell, cell_offsets, ...)

generate_graph(→ dict)

Generate a graph representation from atomic structure data.

Module Contents#

core.graph.compute.get_pbc_distances(pos, edge_index, cell, cell_offsets, neighbors, return_offsets: bool = False, return_distance_vec: bool = False)#
core.graph.compute.generate_graph(data: dict, cutoff: float, max_neighbors: int, enforce_max_neighbors_strictly: bool, radius_pbc_version: int, pbc: torch.Tensor) → dict#

Generate a graph representation from atomic structure data.

Parameters:

  • data (dict) –

    A dictionary containing a batch of molecular structures. It should have the following keys:

    • ’pos’ (torch.Tensor): Positions of the atoms.

    • ’cell’ (torch.Tensor): Cell vectors of the molecular structures.

    • ’natoms’ (torch.Tensor): Number of atoms in each molecular structure.

  • cutoff (float) – The maximum distance between atoms to consider them as neighbors.

  • max_neighbors (int) – The maximum number of neighbors to consider for each atom.

  • enforce_max_neighbors_strictly (bool) – Whether to strictly enforce the maximum number of neighbors.

  • radius_pbc_version – the version of radius_pbc impl

  • pbc (list[bool]) – The periodic boundary conditions in 3 dimensions, defaults to [True,True,True] for 3D pbc

Returns:

A dictionary containing the generated graph with the following keys:

  • ’edge_index’ (torch.Tensor): Indices of the edges in the graph.

  • ’edge_distance’ (torch.Tensor): Distances between the atoms connected by the edges.

  • ’edge_distance_vec’ (torch.Tensor): Vectors representing the distances between the atoms connected by the edges.

  • ’cell_offsets’ (torch.Tensor): Offsets of the cell vectors for each edge.

  • ’offset_distances’ (torch.Tensor): Distances between the atoms connected by the edges, including the cell offsets.

  • ’neighbors’ (torch.Tensor): Number of neighbors for each atom.

Return type:

dict

Contents