torchmdnet.priors

Classes

`Atomref`([max_z, dataset, trainable, enable])	Atomref prior model.
`LearnableAtomref`([max_z])	LearnableAtomref prior model.
`D2`(cutoff_distance, max_num_neighbors[, ...])	Dispersive correction term as used in DFT-D2.
`ZBL`(cutoff_distance, max_num_neighbors[, ...])	Implements the Ziegler-Biersack-Littmark (ZBL) potential for screened nuclear repulsion.
`Coulomb`(lower_switch_distance, ...[, ...])	This class implements a Coulomb potential, scaled by a cosine switching function to reduce its effect at short distances.

class torchmdnet.priors.Atomref(max_z=None, dataset=None, trainable=False, enable=True)[source]

Bases: BasePrior

Atomref prior model.

This prior model is used to add atomic reference values to the input features. The atomic reference values are stored in an embedding layer and are added to the input features as:

\[\begin{split}x' = x + \\textrm{atomref}(z)\end{split}\]

where \(x\) is the input feature tensor, \(z\) is the atomic number tensor, and \(\\textrm{atomref}\) is the embedding layer. The atomic reference values are stored in the embedding layer and can be trainable.

When using this in combination with some dataset, the dataset class must implement the function get_atomref, which returns the atomic reference values as a tensor.

Parameters:

max_z (int, optional) – Maximum atomic number to consider. If dataset is not None, this argument is ignored.
dataset (torch_geometric.data.Dataset, optional) – A dataset from which to extract the atomref values.
trainable (bool, optional) – If False, the atomref values are not trainable. (default: False)
enable (bool, optional) – If False, the prior is disabled. This is useful if you want to add the reference energies only during inference (or training) (default: True)

reset_parameters()[source]

get_init_args()[source]

pre_reduce(x, z, pos, batch=None, extra_args=None)[source]

Applies the stored atomref to the input as:

\[x' = x + \textrm{atomref}(z)\]

Note

The atomref operation is an embedding lookup that can be trainable if the trainable argument is set to False.

Note

This call becomes a no-op if the enable argument is set to False.

Parameters:

x (Tensor) – Input feature tensor.
z (Tensor) – Atomic number tensor.
pos (Tensor) – Atomic positions tensor. Unused.
batch (Tensor, optional) – Batch tensor. Unused. (default: None).
extra_args (Dict[str, Tensor], optional) – Extra arguments. Unused. (default: None)

class torchmdnet.priors.LearnableAtomref(max_z=None)[source]

Bases: Atomref

LearnableAtomref prior model.

This prior model is used to add learned atomic reference values to the input features. The atomic reference values are learned as an embedding layer and are added to the input features as:

\[\begin{split}x' = x + \\textrm{atomref}(z)\end{split}\]

where \(x\) is the input feature tensor, \(z\) is the atomic number tensor, and \(\\textrm{atomref}\) is the embedding layer.

Parameters:: max_z (int, optional) – Maximum atomic number to consider.

class torchmdnet.priors.D2(cutoff_distance, max_num_neighbors, atomic_number=None, distance_scale=None, energy_scale=None, dataset=None, dtype=torch.float32)[source]

Bases: BasePrior

Dispersive correction term as used in DFT-D2.

Reference:: Grimme, Stefan. “Semiempirical GGA-type density functional constructed with a long‐range dispersion correction.” Journal of computational chemistry 27.15 (2006): 1787-1799. Available at: https://onlinelibrary.wiley.com/doi/10.1002/jcc.20495

Parameters:

cutoff_distance (float) – Distance cutoff for the correction term.
max_num_neighbors (int) – Maximum number of neighbors to consider.
atomic_number (list of int, optional) – Map of atom types to atomic numbers. If None, use dataset.atomic_numbers.
position_scale (float, optional) – Factor to convert positions stored in the dataset to meters (m). If None (default), use dataset.position_scale.
energy_scale (float, optional) – Factor to convert energies stored in the dataset to Joules (J). Note: not J/mol. If None (default), use dataset.energy_scale.
dataset (Dataset, optional) – Dataset object. If None, atomic_number, position_scale, and energy_scale must be explicitly set.

Examples

>>> from torchmdnet.priors import D2
>>> prior = D2(
        cutoff_distance=10.0,  # Å
        max_num_neighbors=128,
        atomic_number=list(range(100)),
        position_scale=1e-10,  # Å --> m
        energy_scale=4.35974e-18,  # Hartree --> J
    )

reset_parameters()[source]

get_init_args()[source]

post_reduce(y, z, pos, batch, box=None, extra_args=None)[source]

class torchmdnet.priors.ZBL(cutoff_distance, max_num_neighbors, atomic_number=None, distance_scale=None, energy_scale=None, dataset=None)[source]

Bases: BasePrior

Implements the Ziegler-Biersack-Littmark (ZBL) potential for screened nuclear repulsion.

This potential is described in Ziegler, J.F., Biersack, J.P., Littmark, U. “The Stopping and Range of Ions in Solids.” (1985), specifically in equations 9 and 10 on page 147. It is an empirical potential effectively describing the repulsion between atoms at very short distances.

Reference:: Available at: https://doi.org/10.1007/978-3-642-68779-2_5

Parameters:

atomic_number (torch.Tensor, optional) – A 1D tensor of length max_z. atomic_number[z] is the atomic number of atoms with atom type z. If None, use dataset.atomic_number.
distance_scale (float, optional) – Factor to multiply with coordinates stored in the dataset to convert them to meters. If None, use dataset.distance_scale.
energy_scale (float, optional) – Factor to multiply with energies stored in the dataset to convert them to Joules (not J/mol). If None, use dataset.energy_scale.
dataset (Dataset, optional) – Dataset object. If None, atomic_number, distance_scale, and energy_scale must be explicitly set.

get_init_args()[source]

reset_parameters()[source]

post_reduce(y, z, pos, batch, box=None, extra_args=None)[source]

class torchmdnet.priors.Coulomb(lower_switch_distance, upper_switch_distance, max_num_neighbors, distance_scale=None, energy_scale=None, box_vecs=None, dataset=None)[source]

Bases: BasePrior

This class implements a Coulomb potential, scaled by a cosine switching function to reduce its effect at short distances.

Parameters:

lower_switch_distance (float) – distance below which the interaction strength is zero.
upper_switch_distance (float) – distance above which the interaction has full strength
max_num_neighbors (int) – Maximum number of neighbors per atom allowed.
distance_scale (float, optional) – Factor to multiply with coordinates in the dataset to convert them to meters.
energy_scale (float, optional) – Factor to multiply with energies in the dataset to convert them to Joules (not J/mol).
box_vecs (torch.Tensor, optional) – Initial box vectors for periodic boundary conditions. If None, no periodic boundary conditions are used.
dataset (Dataset) – Dataset object.

Notes

The Dataset used with this class must include a partial_charges field for each sample, and provide distance_scale and energy_scale attributes if they are not explicitly passed as arguments.

get_init_args()[source]

reset_parameters()[source]

post_reduce(y, z, pos, batch, box=None, extra_args=None)[source]

Compute the Coulomb energy for each sample in a batch.

Parameters:

y (torch.Tensor) – Tensor of shape (batch_size, 1) containing the energies of each sample in the batch.
z (torch.Tensor) – Tensor of shape (num_atoms,) containing the atom types for each atom in the batch.
pos (torch.Tensor) – Tensor of shape (num_atoms, 3) containing the positions of each atom in the batch.
batch (torch.Tensor) – Tensor of shape (num_atoms,) containing the batch index for each atom in the batch.
box (torch.Tensor, optional) – Tensor of shape (3, 3) containing the box vectors for the batch. If None, use the initial box vectors.
extra_args (dict, optional) – Dictionary of extra arguments. Must contain a partial_charges field.

Returns:

Tensor of shape (batch_size, 1) containing the energies of each sample in the batch.

Return type:

torch.Tensor