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 of 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])}

See also

BaseAtomFeaturizer, CanonicalAtomFeaturizer, WeaveAtomFeaturizer, AttentiveFPAtomFeaturizer, PAGTNAtomFeaturizer

__init__(atomic_number_types=None, chiral_types=None)[source]

Initialize self. See help(type(self)) for accurate signature.

Methods

__init__([atomic_number_types, chiral_types])

Initialize self.