dgllife.utils.BaseAtomFeaturizer

class dgllife.utils.BaseAtomFeaturizer(featurizer_funcs, feat_sizes=None)[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 and a node i in the graph corresponds to exactly atom i in the molecule.

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 dgllife.utils 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.]])}
>>> # Get feature size for atom mass
>>> print(atom_featurizer.feat_size('mass'))
1
>>> # Get feature size for atom degree
>>> print(atom_featurizer.feat_size('degree'))
11

See also

CanonicalAtomFeaturizer, WeaveAtomFeaturizer, PretrainAtomFeaturizer, AttentiveFPAtomFeaturizer, PAGTNAtomFeaturizer

__init__(featurizer_funcs, feat_sizes=None)[source]

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

Methods

__init__(featurizer_funcs[, feat_sizes])

Initialize self.

feat_size([feat_name])

Get the feature size for feat_name.