MolecularDiffusion.modules.layers.atomic_data¶

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

modified from troch_geometric Data class

Attributes¶

`AseAtomsAdaptor`
`IndexType`

Classes¶

AtomicData

Functions¶

`atomicdata_list_to_batch`(→ AtomicData)	all data points must be single graphs and have the same set of keys.
`get_neighbors_pymatgen`(atoms, cutoff, max_neigh)	Preforms nearest neighbor search and returns edge index, distances,
`reshape_features`(c_index, n_index, n_distance, offsets)	Stack center and neighbor index and reshapes distances,
`size_repr`(→ str)
`tensor_or_int_to_tensor`(x[, dtype])

Module Contents¶

class MolecularDiffusion.modules.layers.atomic_data.AtomicData(pos: torch.Tensor, atomic_numbers: torch.Tensor, cell: torch.Tensor, pbc: torch.Tensor, natoms: torch.Tensor, edge_index: torch.Tensor, cell_offsets: torch.Tensor, nedges: torch.Tensor, charge: torch.Tensor, spin: torch.Tensor, fixed: torch.Tensor, tags: torch.Tensor, energy: torch.Tensor | None = None, forces: torch.Tensor | None = None, stress: torch.Tensor | None = None, batch: torch.Tensor | None = None, sid: list[str] | None = None, dataset: list[str] | str | None = None)¶

apply(func)¶: Applies the function func to all tensor attributes

assign_batch_stats(slices, cumsum, cat_dims, natoms_list)¶

batch_to_atomicdata_list() → list[AtomicData]¶: Reconstructs the list of torch_geometric.data.Data objects from the batch object. The batch object must have been created via from_data_list() in order to be able to reconstruct the initial objects.

clone()¶: Performs a deep-copy of the data object.

contiguous()¶: Ensures a contiguous memory layout for all tensor attributes

cpu()¶: Copies all tensor attributes to CPU memory.

cuda(device=None, non_blocking=False)¶: Copies all tensor attributes to GPU memory.

classmethod from_ase(input_atoms: ase.Atoms, r_edges: bool = False, radius: float = 6.0, max_neigh: int | None = None, sid: str | None = None, molecule_cell_size: float | None = None, r_energy: bool = True, r_forces: bool = True, r_stress: bool = True, r_data_keys: list[str] | None = None, task_name: str | None = None) → AtomicData¶

classmethod from_dict(dictionary)¶: Creates a data object from a python dictionary.

get(key, default)¶

get_batch_stats()¶

get_example(idx: int) → AtomicData¶: Reconstructs the AtomicData object at index idx from a batched AtomicData object.

index_select(idx: IndexType) → list[AtomicData]¶

keys()¶

to(device, **kwargs)¶: Performs tensor dtype and/or device conversion for all tensor attributes

to_ase() → list[ase.Atoms]¶

to_ase_single() → ase.Atoms¶

to_dict()¶

validate()¶

values()¶

atomic_numbers¶

cell¶

cell_offsets¶

charge¶

edge_index¶

fixed¶

natoms¶

nedges¶

property num_edges: int¶: Returns the number of edges in the graph.

property num_graphs: int¶: Returns the number of graphs in the batch.

property num_nodes: int¶: Returns or sets the number of nodes in the graph.

pbc¶

pos¶

sid¶

spin¶

tags¶

property task_name¶

MolecularDiffusion.modules.layers.atomic_data.atomicdata_list_to_batch(data_list: list[AtomicData], exclude_keys: list | None = None) → AtomicData¶: all data points must be single graphs and have the same set of keys. TODO: exclude keys?

MolecularDiffusion.modules.layers.atomic_data.get_neighbors_pymatgen(atoms: ase.Atoms, cutoff, max_neigh)¶: Preforms nearest neighbor search and returns edge index, distances, and cell offsets

MolecularDiffusion.modules.layers.atomic_data.reshape_features(c_index: numpy.ndarray, n_index: numpy.ndarray, n_distance: numpy.ndarray, offsets: numpy.ndarray)¶: Stack center and neighbor index and reshapes distances, takes in np.arrays and returns torch tensors

MolecularDiffusion.modules.layers.atomic_data.size_repr(key: str, item: torch.Tensor, indent=0) → str¶

MolecularDiffusion.modules.layers.atomic_data.tensor_or_int_to_tensor(x, dtype=torch.int)¶

MolecularDiffusion.modules.layers.atomic_data.AseAtomsAdaptor = None¶

MolecularDiffusion.modules.layers.atomic_data.IndexType¶