dgllife.utils.PretrainAtomFeaturizer¶

class dgllife.utils.PretrainAtomFeaturizer(atomic_number_types=None, chiral_types=None)[source]¶

AtomFeaturizer in Strategies for Pre-training Graph Neural Networks.

The atom featurization performed in Strategies for Pre-training Graph Neural Networks, which considers:

atomic number
chirality

We assume the resulting DGLGraph will not contain any virtual nodes.

Parameters

atomic_number_types (list of int or None) – Atomic number types to consider for one-hot encoding. If None, we will use a default choice of 1-118.
chiral_types (list of Chem.rdchem.ChiralType or None) – Atom chirality to consider for one-hot encoding. If None, we will use a default choice, including Chem.rdchem.ChiralType.CHI_UNSPECIFIED, Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CW, Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CCW, Chem.rdchem.ChiralType.CHI_OTHER.

Examples

>>> from rdkit import Chem
>>> from dgllife.utils import PretrainAtomFeaturizer

>>> mol = Chem.MolFromSmiles('CCO')
>>> atom_featurizer = PretrainAtomFeaturizer()
>>> atom_featurizer(mol)
{'atomic_number': tensor([5, 5, 7]), 'chirality_type': tensor([0, 0, 0])}