torchdrug.data

Data Structures

Datasets

Data Processing

Feature Functions

Element Constants

Element constants are provided for convenient manipulation of atom types. The atomic numbers can be accessed by uppercased element names at the root of the package. For example, we can get the carbon scaffold of a molecule with the following code.

import torchdrug as td
from torchdrug import data

smiles = "CC1=C(C=C(C=C1[N+](=O)[O-])[N+](=O)[O-])[N+](=O)[O-]"
mol = data.Molecule.from_smiles(smiles)
scaffold = mol.subgraph(mol.atom_type == td.CARBON)
mol.visualize()
scaffold.visualize()

There are also 2 constant arrays that map atomic numbers to element names. td.ATOM_NAME[i] returns the full name, while td.ATOM_SYMBOL[i] returns the abbreviated chemical symbol for atomic number i.

For a full list of elements, please refer to the perodic table.