SO(3) Lie Group Operations

The module nnp.so3 contains tools to rotate point clouds in 3D space.

nnp.so3.rotate_along(axis: torch.Tensor) → torch.Tensor

Compute group elements of rotating along an axis passing origin.

Parameters:axis – a vector (x, y, z) whose direction specifies the axis of the rotation, length specifies the radius to rotate, and sign specifies clockwise or anti-clockwise. Returns:the rotational matrix \(\exp{\left(\theta W\right)}\).

Molecular Dynamics

The module nnp.md provide tools to run molecular dynamics with a potential defined by PyTorch.

class nnp.md.Calculator(func: Callable[[Sequence[str], torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor], overwrite: bool = False)

ASE Calculator that wraps a neural network potential

Parameters:

• func (callable) – A fucntion that .
• overwrite (bool) – After wrapping atoms into central box, whether to replace the original positions stored in ase.Atoms object with the wrapped positions.

calculate(atoms=None, properties=['energy'], system_changes=['positions', 'numbers', 'cell', 'pbc', 'initial_charges', 'initial_magmoms'])

Do the calculation.

properties: list of str

List of what needs to be calculated. Can be any combination of ‘energy’, ‘forces’, ‘stress’, ‘dipole’, ‘charges’, ‘magmom’ and ‘magmoms’.

system_changes: list of str

List of what has changed since last calculation. Can be any combination of these six: ‘positions’, ‘numbers’, ‘cell’, ‘pbc’, ‘initial_charges’ and ‘initial_magmoms’.

Subclasses need to implement this, but can ignore properties and system_changes if they want. Calculated properties should be inserted into results dictionary like shown in this dummy example:

self.results = {'energy': 0.0,
                'forces': np.zeros((len(atoms), 3)),
                'stress': np.zeros(6),
                'dipole': np.zeros(3),
                'charges': np.zeros(len(atoms)),
                'magmom': 0.0,
                'magmoms': np.zeros(len(atoms))}

The subclass implementation should first call this implementation to set the atoms attribute.

Periodic Boundary Conditions

The module nnp.pbc contains tools to deal with periodic boundary conditions.

nnp.pbc.map2central(cell: torch.Tensor, coordinates: torch.Tensor, pbc: torch.Tensor) → torch.Tensor

Map atoms outside the unit cell into the cell using PBC.

Parameters:

• cell –

  tensor of shape (3, 3) of the three vectors defining unit cell:

  tensor([[x1, y1, z1],
          [x2, y2, z2],
          [x3, y3, z3]])
  

• coordinates – Tensor of shape (atoms, 3) or (molecules, atoms, 3).
• pbc – boolean vector of size 3 storing if pbc is enabled for that direction.

Returns:

coordinates of atoms mapped back to unit cell.

nnp.pbc.num_repeats(cell: torch.Tensor, pbc: torch.Tensor, cutoff: float) → torch.Tensor

Compute the number of repeats required along each cell vector to make the original cell and repeated cells together form a large enough box so that for each atom in the original cell, all its neighbor atoms within the given cutoff distance are contained in the big box.

Parameters:

• cell –

  tensor of shape (3, 3) of the three vectors defining unit cell:

  tensor([[x1, y1, z1],
          [x2, y2, z2],
          [x3, y3, z3]])
  

• pbc – boolean vector of size 3 storing if pbc is enabled for that direction.
• cutoff – the cutoff inside which atoms are considered as pairs

Returns:

numbers of repeats required to make the repeated box contains all the neighbors of atoms in the ceter cell

Vibrational Analysis

The module nnp.vib contains tools to compute analytical hessian and do vibrational analysis.

class nnp.vib.FreqsModes(angular_frequencies, modes)

property angular_frequencies

Alias for field number 0

property modes

Alias for field number 1

nnp.vib.hessian(coordinates: torch.Tensor, energies: Optional[torch.Tensor] = None, forces: Optional[torch.Tensor] = None) → torch.Tensor

Compute analytical hessian from the energy graph or force graph.

Parameters:

• coordinates – Tensor of shape (molecules, atoms, 3) or (atoms, 3)
• energies – Tensor of shape (molecules,), or scalar, if specified, then forces must be None. This energies must be computed from coordinates in a graph.
• forces – Tensor of shape (molecules, atoms, 3) or (atoms, 3), if specified, then energies must be None. This forces must be computed from coordinates in a graph.

Returns:

Tensor of shape (molecules, 3 * atoms, 3 * atoms) or (3 * atoms, 3 * atoms)

nnp.vib.vibrational_analysis(masses: torch.Tensor, hessian: torch.Tensor) → nnp.vib.FreqsModes

Computing the vibrational wavenumbers from hessian.

Parameters:

• masses – Tensor of shape (molecules, atoms) or (atoms,).
• hessian – Tensor of shape (molecules, 3 * atoms, 3 * atoms) or (3 * atoms, 3 * atoms).

Returns:

A namedtuple (angular_frequencies, modes) where

angular_frequencies:

Tensor of shape (molecules, 3 * atoms) or (3 * atoms,)

modes:

Tensor of shape (molecules, modes, atoms, 3) or (modes, atoms, 3) where modes = 3 * atoms is the number of normal modes.