faenet.utils

Module Contents

Classes

GaussianSmearing	Smears a distance distribution by a Gaussian function.
RandomReflect	Reflect node positions around a specific axis (x, y, x=y) or the origin.
RandomRotate	Rotates node positions around a specific axis by a randomly sampled

Functions

base_preprocess(data[, cutoff, max_num_neighbors])	Preprocess datapoint: create a cutoff graph,
get_pbc_distances(pos, edge_index, cell, cell_offsets, ...)	Compute distances between atoms with periodic boundary conditions
pbc_preprocess(data[, cutoff, max_num_neighbors])	Preprocess datapoint using periodic boundary conditions
swish(x)	Swish activation function

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 datapoint: create a cutoff graph,

compute distances and relative positions.

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 datapoint using periodic boundary conditions

to improve the existing graph.

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)

Swish activation function