torchani.potentials

Callable objects that can be attached to TorchANI models.

Potentials calculate explicit analytical energy expressions that modify the final output of the models, usually to correct it in some regions of chemical space where they don’t behave in a physically appropriate manner.

New potentials can be added to the TorchANI library by simple subclassing. For a tutorial on how to do this please refer to the corresponding documentation in examples.

Potential is a torch Module. It can represent many-body potentials. PairPotential is a more specific subclass, which represents pair potentials. Subclasses of Potential should override compute, and subclasses of PairPotential should override pair_energies.

Many potentials correspond to functions implemented in specific scientific articles or books. If you use any of these potentials in your work, please include the corresponding citations.

Classes

Potential	Base class for all atomic potentials
DummyPotential
BasePairPotential	General base class for all pairwise potentials
PairPotential	User friendly, simple class for pairwise potentials
RepulsionXTB	Calculates the xTB repulsion energy terms for a given molecule
TwoBodyDispersionD3	Calculates the DFT-D3 dispersion corrections
NNPotential
SeparateChargesNNPotential
MergedChargesNNPotential
RepulsionZBL
RepulsionLJ
DispersionLJ
LennardJones
FixedCoulomb
FixedMNOK

class torchani.potentials.Potential(symbols, *, cutoff=inf)

Base class for all atomic potentials

Potentials may be many-body (3-body, 4-body, …) potentials or 2-body (pair) potentials. Subclasses must implement compute and may override __init__.

compute_from_neighbors(elem_idxs, coords, neighbors, charge=0, atomic=False, ensemble_values=False, ghost_flags=None)

Compute the energies associated with the potential

If the potential is an ensemble of multiple models, ensemble_values=True should return the individual values of the models, in the first dimension (submodels, ...). Otherwise it should disregard ensemble_values.

Must be implemented by subclasses

forward(species, coords, cell=None, pbc=None, atomic=False, ensemble_values=False, atomic_nums_input=True)

Outputs energy, as calculated by the potential

Output shape depends on the value of atomic, it is either (molecs, atoms) or (molecs,)

class torchani.potentials.DummyPotential(symbols, *, cutoff=inf)

compute_from_neighbors(elem_idxs, coords, neighbors, charge=0, atomic=False, ensemble_values=False, ghost_flags=None)

Compute the energies associated with the potential

If the potential is an ensemble of multiple models, ensemble_values=True should return the individual values of the models, in the first dimension (submodels, ...). Otherwise it should disregard ensemble_values.

Must be implemented by subclasses

class torchani.potentials.BasePairPotential(symbols, *, cutoff=inf, cutoff_fn='smooth')

General base class for all pairwise potentials

Subclasses must implement pair_energies, and override init using this template:

def __init__(
    symbols,
    ..., # User args go here
    cutoff: float=math.inf,
    cutoff_fn="smooth",
    ..., # User kwargs go here
)
    super().__init__(symbols, cutoff, cutoff_fn)
    ... # User code goes here

compute_from_neighbors(elem_idxs, coords, neighbors, charge=0, atomic=False, ensemble_values=False, ghost_flags=None)

Compute the energies associated with the potential

If the potential is an ensemble of multiple models, ensemble_values=True should return the individual values of the models, in the first dimension (submodels, ...). Otherwise it should disregard ensemble_values.

Must be implemented by subclasses

pair_energies(elem_idxs, neighbors)

Return energy of all pairs of neighbors

Returns a Tensor of energies, of shape (‘pairs’,) where ‘pairs’ is the number of neighbor pairs.

static symm(x)

Takes an NxN tensor with 0 lower triangle and returns it symmetrized

to_pair_values(x, elem_idxs)

Returns values of elem pairs from a NxN tensor with 0 low triangle

class torchani.potentials.PairPotential(symbols, *, trainable=(), cutoff=inf, cutoff_fn='smooth', **kwargs)

User friendly, simple class for pairwise potentials

Subclasses must implement pair_energies and, if they use any parameters or buffers, specify three list of strings:

• 'tensors': Vectors (all with the same len) or scalars

• 'elem_tensors': With shape (num-sym,)

• 'pair_elem_tensors': With shape (num-sym * (num-sym + 1) / 2,)

Usage is better understood by an example:

from torchani.potentials import PairPotential


class Square(PairPotential):
    tensors = ['bias']  # Vectors (all with the same len) or scalars
    pair_elem_tensors = ["k", "eq"]  # shape (num-sym * (num-sym + 1) / 2)

    def pair_energies(self, elem_idxs, neighbors):
        elem_pairs = elem_idxs.view(-1)[neighbors.indices]
        eq = self.to_pair_values(self.eq, elem_pairs)
        k = self.to_pair_values(self.k, elem_pairs)
        return self.bias + k / 2 * (neighbors.distances - eq) ** 2


#  Order for the pair elem tensors is HH, HC, HO, CC, CO, ...
#  Values for demonstration purpose only
k = (1.,) * (3 * (3 + 1) // 2)
eq = (1.5,) * (3 * (3 + 1) // 2)
pot = Square(symbols=("H", "C", "O"), k=k, eq=eq, bias=0.1)

# Or if the constants are trainable:
pot = Square(symbols=("H", "C", "O"), k=k, bias=0.1, eq=eq, trainable="k")

class torchani.potentials.RepulsionXTB(symbols, krep_hydrogen=1.0, krep=1.5, alpha=(), yeff=(), *, cutoff=inf, cutoff_fn='smooth')

Calculates the xTB repulsion energy terms for a given molecule

Potential used is as in work by Grimme: https://pubs.acs.org/doi/10.1021/acs.jctc.8b01176

By default alpha, yeff and krep parameters are taken from Grimme et. al.

krep_hydrogen is only used for H-H repulsive interaction. All other interactions use krep

pair_energies(element_idxs, neighbors)

Return energy of all pairs of neighbors

Returns a Tensor of energies, of shape (‘pairs’,) where ‘pairs’ is the number of neighbor pairs.

class torchani.potentials.TwoBodyDispersionD3(symbols, s6, s8, damp_a1, damp_a2, sqrt_empirical_charge=(), covalent_radii=(), *, cutoff_fn='smooth', cutoff=inf)

Calculates the DFT-D3 dispersion corrections

Only calculates the 2-body part of the dispersion corrections. Requires a damping function for the order-6 and order-8 potential terms.

pair_energies(element_idxs, neighbors)

Return energy of all pairs of neighbors

Returns a Tensor of energies, of shape (‘pairs’,) where ‘pairs’ is the number of neighbor pairs.

class torchani.potentials.NNPotential(aev_computer, neural_networks)

compute_from_neighbors(elem_idxs, coords, neighbors, charge=0, atomic=False, ensemble_values=False, ghost_flags=None)

Compute the energies associated with the potential

If the potential is an ensemble of multiple models, ensemble_values=True should return the individual values of the models, in the first dimension (submodels, ...). Otherwise it should disregard ensemble_values.

Must be implemented by subclasses

class torchani.potentials.SeparateChargesNNPotential(aev_computer, neural_networks, charge_networks, charge_normalizer=None)

compute_from_neighbors(elem_idxs, coords, neighbors, charge=0, atomic=False, ensemble_values=False, ghost_flags=None)

Compute the energies associated with the potential

If the potential is an ensemble of multiple models, ensemble_values=True should return the individual values of the models, in the first dimension (submodels, ...). Otherwise it should disregard ensemble_values.

Must be implemented by subclasses

class torchani.potentials.MergedChargesNNPotential(aev_computer, neural_networks, charge_normalizer=None)

compute_from_neighbors(elem_idxs, coords, neighbors, charge=0, atomic=False, ensemble_values=False, ghost_flags=None)

Compute the energies associated with the potential

If the potential is an ensemble of multiple models, ensemble_values=True should return the individual values of the models, in the first dimension (submodels, ...). Otherwise it should disregard ensemble_values.

Must be implemented by subclasses

class torchani.potentials.RepulsionZBL(symbols, k=0.8853, screen_coeffs=(), screen_exponents=(), eff_exponent=0.23, eff_atomic_nums=(), trainable=(), *, cutoff=inf, cutoff_fn='smooth')

pair_energies(elem_idxs, neighbors)

Return energy of all pairs of neighbors

Returns a Tensor of energies, of shape (‘pairs’,) where ‘pairs’ is the number of neighbor pairs.

class torchani.potentials.RepulsionLJ(symbols, eps=(), sigma=(), *, trainable=(), cutoff=inf, cutoff_fn='smooth')

pair_energies(elem_idxs, neighbors)

Return energy of all pairs of neighbors

Returns a Tensor of energies, of shape (‘pairs’,) where ‘pairs’ is the number of neighbor pairs.

class torchani.potentials.DispersionLJ(symbols, eps=(), sigma=(), *, trainable=(), cutoff=inf, cutoff_fn='smooth')

pair_energies(elem_idxs, neighbors)

Return energy of all pairs of neighbors

Returns a Tensor of energies, of shape (‘pairs’,) where ‘pairs’ is the number of neighbor pairs.

class torchani.potentials.LennardJones(symbols, eps=(), sigma=(), *, trainable=(), cutoff=inf, cutoff_fn='smooth')

pair_energies(elem_idxs, neighbors)

Return energy of all pairs of neighbors

Returns a Tensor of energies, of shape (‘pairs’,) where ‘pairs’ is the number of neighbor pairs.

class torchani.potentials.FixedCoulomb(symbols, dielectric=1.0, charges=(), *, cutoff=inf, cutoff_fn='smooth')

pair_energies(elem_idxs, neighbors)

Return energy of all pairs of neighbors

Returns a Tensor of energies, of shape (‘pairs’,) where ‘pairs’ is the number of neighbor pairs.

class torchani.potentials.FixedMNOK(symbols, dielectric=1.0, charges=(), eta=(), *, cutoff=inf, cutoff_fn='smooth')

pair_energies(elem_idxs, neighbors)

Return energy of all pairs of neighbors

Returns a Tensor of energies, of shape (‘pairs’,) where ‘pairs’ is the number of neighbor pairs.