dgllife.utils.PretrainBondFeaturizer

class dgllife.utils.PretrainBondFeaturizer(bond_types=None, bond_direction_types=None, self_loop=True)[source]

BondFeaturizer in Strategies for Pre-training Graph Neural Networks.

The bond featurization performed in Strategies for Pre-training Graph Neural Networks, which considers:

  • bond type

  • bond direction

Parameters

  • bond_types (list of Chem.rdchem.BondType or None) – Bond types to consider. Default to Chem.rdchem.BondType.SINGLE, Chem.rdchem.BondType.DOUBLE, Chem.rdchem.BondType.TRIPLE, Chem.rdchem.BondType.AROMATIC.

  • bond_direction_types (list of Chem.rdchem.BondDir or None) – Bond directions to consider. Default to Chem.rdchem.BondDir.NONE, Chem.rdchem.BondDir.ENDUPRIGHT, Chem.rdchem.BondDir.ENDDOWNRIGHT.

  • self_loop (bool) – Whether self loops will be added. Default to True.

Examples

>>> from dgllife.utils import PretrainBondFeaturizer
>>> from rdkit import Chem

>>> mol = Chem.MolFromSmiles('CO')
>>> bond_featurizer = PretrainBondFeaturizer()
>>> bond_featurizer(mol)
{'bond_type': tensor([0, 0, 4, 4]),
 'bond_direction_type': tensor([0, 0, 0, 0])}
__init__(bond_types=None, bond_direction_types=None, self_loop=True)[source]

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

Methods

__init__([bond_types, bond_direction_types, …])

Initialize self.