faenet.frame_averaging

Module Contents

Functions

check_constraints(eigenval, eigenvec[, dim])

Check requirements for frame averaging are satisfied

compute_frames(eigenvec, pos, cell[, fa_method, ...])

Compute all frames for a given graph.

data_augmentation(g[, d])

Data augmentation where we add randomly rotated graphs

frame_averaging_2D(pos[, cell, fa_method, check])

Computes new positions for the graph atoms,

frame_averaging_3D(pos[, cell, fa_method, check])

Computes new positions for the graph atoms using PCA
faenet.frame_averaging.check_constraints(eigenval, eigenvec, dim=3)

Check requirements for frame averaging are satisfied

Parameters:

  • eigenval (tensor) – eigenvalues

  • eigenvec (tensor) – eigenvectors

  • dim (int) – 2D or 3D frame averaging

faenet.frame_averaging.compute_frames(eigenvec, pos, cell, fa_method='stochastic', pos_3D=None, det_index=0)

Compute all frames for a given graph.

Parameters:

  • eigenvec (tensor) – eigenvectors matrix

  • pos (tensor) – centered position vector

  • cell (tensor) – cell direction (dxd)

  • fa_method (str) – the Frame Averaging (FA) inspired technique chosen to select frames: stochastic-FA (stochastic), deterministic-FA (det), Full-FA (all) or SE(3)-FA (se3).

  • pos_3D – for 2D FA, pass atoms’ 3rd position coordinate.

Returns:

3D position tensors of projected representation

Return type:

list

faenet.frame_averaging.data_augmentation(g, d=3, *args)

Data augmentation where we add randomly rotated graphs in the dataloader transform.

Parameters:

  • g (data.Data) – single graph

  • d (int) – dimension of the DA rotation (2D around z-axis or 3D)

  • rotation (str, optional) – around which axis do we rotate it. Defaults to ‘z’.

Returns:

rotated graph

Return type:

(data.Data)

faenet.frame_averaging.frame_averaging_2D(pos, cell=None, fa_method='stochastic', check=False)

Computes new positions for the graph atoms, based on a frame averaging building on PCA.

Parameters:

  • pos (tensor) – positions of atoms in the graph

  • cell (tensor) – unit cell of the graph. None if no pbc.

  • fa_method (str) – FA method used (stochastic, det, all, se3)

  • check (bool) – check if constraints are satisfied. Default: False.

Returns:

updated atom positions tensor: updated unit cell tensor: the rotation matrix used (PCA)

Return type:

tensor

faenet.frame_averaging.frame_averaging_3D(pos, cell=None, fa_method='stochastic', check=False)

Computes new positions for the graph atoms using PCA

Parameters:

  • pos (tensor) – positions of atoms in the graph

  • cell (tensor) – unit cell of the graph. None if no pbc.

  • fa_method (str) – FA method used (stochastic, det, all, se3-all, se3-det, se3-stochastic)

  • check (bool) – check if constraints are satisfied. Default: False.

Returns:

updated atom positions tensor: updated unit cell tensor: the rotation matrix used (PCA)

Return type:

tensor