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

__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`.