dgl.data.chem.BaseAtomFeaturizer

class dgl.data.chem.BaseAtomFeaturizer(**kwargs)[source]

An abstract class for atom featurizers.

Loop over all atoms in a molecule and featurize them with the featurizer_funcs.

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

Parameters:
  • featurizer_funcs (dict) – Mapping feature name to the featurization function. Each function is of signature func(rdkit.Chem.rdchem.Atom) -> list or 1D numpy array.
  • feat_sizes (dict) – Mapping feature name to the size of the corresponding feature. If None, they will be computed when needed. Default: None.

Examples

>>> from dgl.data.chem import BaseAtomFeaturizer, atom_mass, atom_degree_one_hot
>>> from rdkit import Chem
>>> mol = Chem.MolFromSmiles('CCO')
>>> atom_featurizer = BaseAtomFeaturizer({'mass': atom_mass, 'degree': atom_degree_one_hot})
>>> atom_featurizer(mol)
{'mass': tensor([[0.1201],
                 [0.1201],
                 [0.1600]]),
 'degree': tensor([[0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
                   [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.],
                   [0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])}
__init__(**kwargs)

Methods

__init__(**kwargs)
feat_size(feat_name) Get the feature size for feat_name.