faenet.utils
Module Contents
Classes
Smears a distance distribution by a Gaussian function. |
|
Reflect node positions around a specific axis (x, y, x=y) or the origin. |
|
Rotates node positions around a specific axis by a randomly sampled |
Functions
|
Preprocess data using a simple cutoff radius |
|
Compute distances between atoms with periodic boundary conditions |
|
Preprocess data using periodic boundary conditions |
|
- class faenet.utils.GaussianSmearing(start=0.0, stop=5.0, num_gaussians=50)
Bases:
torch.nn.Module
Smears a distance distribution by a Gaussian function.
- forward(dist)
- class faenet.utils.RandomReflect
Bases:
object
Reflect node positions around a specific axis (x, y, x=y) or the origin. Take a random reflection type from a list of reflection types.
(type 0: reflect wrt x-axis, type1: wrt y-axis, type2: y=x, type3: origin)
- __call__(data)
- class faenet.utils.RandomRotate(degrees, axes=[0, 1, 2])
Bases:
object
Rotates node positions around a specific axis by a randomly sampled factor within a given interval.
- Parameters:
degrees (tuple or float) – Rotation interval from which the rotation angle is sampled. If degrees is a number instead of a tuple, the interval is given by \([-\mathrm{degrees}, \mathrm{degrees}]\).
axes (int, optional) – The rotation axes. (default: [0, 1, 2])
- __call__(data)
- __repr__()
Return repr(self).
- faenet.utils.base_preprocess(data, cutoff=6.0, max_num_neighbors=40)
Preprocess data using a simple cutoff radius
Args: data (data.Data): data object with specific attributes:
batch (N): index of the graph to which each atom belongs to in this batch
pos (N,3): atom positions
atomic_numbers (N): atomic numbers of each atom in the batch
edge_index (2,E): edge indices, for all graphs of the batch
With B is the batch size, N the number of atoms in the batch (across all graphs), and E the number of edges in the batch. If these attributes are not present, implement your own preprocess function.
cutoff (int): cutoff radius (in Angstrom) max_num_neighbors (int): maximum number of neighbors per node.
- Returns:
(atomic_numbers, batch, sparse adjacency matrix, relative positions, distances)
- Return type:
tuple
- faenet.utils.get_pbc_distances(pos, edge_index, cell, cell_offsets, neighbors, return_offsets=False, return_rel_pos=False)
Compute distances between atoms with periodic boundary conditions
- Parameters:
pos (tensor) – (N, 3) tensor of atomic positions
edge_index (tensor) – (2, E) tensor of edge indices
cell (tensor) – (3, 3) tensor of cell vectors
cell_offsets (tensor) – (N, 3) tensor of cell offsets
neighbors (tensor) – (N, 3) tensor of neighbor indices
return_offsets (bool) – return the offsets
return_rel_pos (bool) – return the relative positions vectors
- Returns:
- dictionary with the updated edge_index, atom distances,
and optionally the offsets and distance vectors.
- Return type:
dict
- faenet.utils.pbc_preprocess(data, cutoff=6.0, max_num_neighbors=40)
Preprocess data using periodic boundary conditions Used for OC20.
- Parameters:
data (data.Data) –
data object with specific attributes. B is the batch size, N the number of atoms in the batch (across all graphs), E the number of edges in the batch.
batch (N): index of the graph to which each atom belongs to in this batch
pos (N,3): atom positions
atomic_numbers (N): atomic numbers of each atom in the batch
cell (B, 3, 3): unit cell containing each graph, for materials.
cell_offsets (E, 3): cell offsets for each edge, for materials
neighbors (B): total number of edges inside each graph.
edge_index (2,E): edge indices, for all graphs of the batch
If these attributes are not present, implement your own preprocess function.
- Returns:
(atomic_numbers, batch, sparse adjacency matrix, relative positions, distances)
- Return type:
tuple
- faenet.utils.swish(x)