dgllife.utils.BaseBondFeaturizer

class dgllife.utils.BaseBondFeaturizer(featurizer_funcs, feat_sizes=None, self_loop=False)[source]

An abstract class for bond featurizers. Loop over all bonds in a molecule and featurize them with the featurizer_funcs. We assume the constructed DGLGraph is a bi-directed graph where the i th bond in the molecule, i.e. mol.GetBondWithIdx(i), corresponds to the (2i)-th and (2i+1)-th edges in the DGLGraph.

We assume the resulting DGLGraph will be created with :func:`smiles_to_bigraph` without self loops.

Parameters
  • featurizer_funcs (dict) – Mapping feature name to the featurization function. Each function is of signature func(rdkit.Chem.rdchem.Bond) -> 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.

  • self_loop (bool) – Whether self loops will be added. Default to False. If True, it will use an additional column of binary values to indicate the identity of self loops in each bond feature. The features of the self loops will be zero except for the additional columns.

Examples

>>> from dgllife.utils import BaseBondFeaturizer, bond_type_one_hot, bond_is_in_ring
>>> from rdkit import Chem
>>> mol = Chem.MolFromSmiles('CCO')
>>> bond_featurizer = BaseBondFeaturizer({'type': bond_type_one_hot, 'ring': bond_is_in_ring})
>>> bond_featurizer(mol)
{'type': tensor([[1., 0., 0., 0.],
                 [1., 0., 0., 0.],
                 [1., 0., 0., 0.],
                 [1., 0., 0., 0.]]),
 'ring': tensor([[0.], [0.], [0.], [0.]])}
>>> # Get feature size
>>> bond_featurizer.feat_size('type')
4
>>> bond_featurizer.feat_size('ring')
1

# Featurization with self loops to add

>>> bond_featurizer = BaseBondFeaturizer(
...                       {'type': bond_type_one_hot, 'ring': bond_is_in_ring},
...                       self_loop=True)
>>> bond_featurizer(mol)
{'type': tensor([[1., 0., 0., 0., 0.],
                 [1., 0., 0., 0., 0.],
                 [1., 0., 0., 0., 0.],
                 [1., 0., 0., 0., 0.],
                 [0., 0., 0., 0., 1.],
                 [0., 0., 0., 0., 1.],
                 [0., 0., 0., 0., 1.]]),
 'ring': tensor([[0., 0.],
                 [0., 0.],
                 [0., 0.],
                 [0., 0.],
                 [0., 1.],
                 [0., 1.],
                 [0., 1.]])}
>>> # Get feature size
>>> bond_featurizer.feat_size('type')
5
>>> bond_featurizer.feat_size('ring')
2
__init__(featurizer_funcs, feat_sizes=None, self_loop=False)[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.