dgllife.utils.PAGTNEdgeFeaturizer

class dgllife.utils.PAGTNEdgeFeaturizer(bond_data_field='e', max_length=5)[source]

The edge featurizer used in PAGTN

PAGTN is introduced in Path-Augmented Graph Transformer Network.

We build a complete graph and the edge features include: * Shortest path between two nodes in terms of bonds. To encode the path,

we encode each bond on the path and concatenate their encodings. The encoding of a bond contains information about the bond type, whether the bond is conjugated and whether the bond is in a ring.

  • One hot encoding of type of rings based on size and aromaticity.

  • One hot encoding of the distance between the nodes.

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

Parameters
  • bond_data_field (str) – Name for storing bond features in DGLGraphs, default to 'e'.

  • max_length (int) – Maximum distance up to which shortest paths must be considered. Paths shorter than max_length will be padded and longer will be truncated, default to 5.

Examples

>>> from dgllife.utils import PAGTNEdgeFeaturizer
>>> from rdkit import Chem
>>> mol = Chem.MolFromSmiles('CCO')
>>> bond_featurizer = PAGTNEdgeFeaturizer(max_length=1)
>>> bond_featurizer(mol)
{'e': tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
              [1., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
              [1., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
              [1., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
              [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
              [1., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
              [1., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
              [1., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0.],
              [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])}
>>> # Get feature size
>>> bond_featurizer.feat_size()
14
__init__(bond_data_field='e', max_length=5)[source]

Methods

__init__([bond_data_field, max_length])

bond_features(mol, path_atoms, ring_info)

Computes the edge features for a given pair of nodes.

feat_size()

Get the feature size.