def get_point_weight(mol, probe_rad=1.4):
"""Form a set of weighted points, with each point representing an atom in a molecule.
Parameters:
mol - rdkit.Chem.rdchem.Mol molecule
probe_rad - probe radius, default of 1.4 Angstrom
Returns the weights of the points (squared (VDW radius + probe radius)).
"""
coor = mol.GetConformer().GetPositions()
tbl = GetPeriodicTable()
pts_num = coor.shape[0]
weights = np.zeros(pts_num)
for i in range(pts_num):
weights[i] = (tbl.GetRvdw(mol.GetAtomWithIdx(int(i)).GetAtomicNum()) +
probe_rad)**2
return weights
def mol_to_dgl(mol):
"""Featurizes an rdkit mol object to a DGL Graph, with node and edge features
Parameters
----------
mol : rdkit mol
Returns
-------
dgl.graph
"""
g = dgl.DGLGraph()
g.add_nodes(mol.GetNumAtoms())
g.set_n_initializer(dgl.init.zero_initializer)
g.set_e_initializer(dgl.init.zero_initializer)
# Atom features
atom_features = []
pd = GetPeriodicTable()
# ComputeGasteigerCharges(mol)
for atom in mol.GetAtoms():
atom_feat = []
atom_type = [0] * len(ATOM_TYPES)
atom_type[ATOM_TYPES.index(atom.GetSymbol())] = 1
chiral = [0] * len(CHIRALITY)
chiral[CHIRALITY.index(atom.GetChiralTag())] = 1
ex_valence = atom.GetExplicitValence()
charge = atom.GetFormalCharge()
hybrid = [0] * len(HYBRIDIZATION)
hybrid[HYBRIDIZATION.index(atom.GetHybridization())] = 1
degree = atom.GetDegree()
valence = atom.GetImplicitValence()
aromatic = int(atom.GetIsAromatic())
ex_hs = atom.GetNumExplicitHs()
im_hs = atom.GetNumImplicitHs()
rad = atom.GetNumRadicalElectrons()
ring = int(atom.IsInRing())
mass = pd.GetAtomicWeight(atom.GetSymbol())
vdw = pd.GetRvdw(atom.GetSymbol())
# pcharge = float(atom.GetProp("_GasteigerCharge"))
atom_feat.extend(atom_type)
atom_feat.extend(chiral)
atom_feat.append(ex_valence)
atom_feat.append(charge)
atom_feat.extend(hybrid)
atom_feat.append(degree)
atom_feat.append(valence)
atom_feat.append(aromatic)
atom_feat.append(ex_hs)
atom_feat.append(im_hs)
atom_feat.append(rad)
atom_feat.append(ring)
atom_feat.append(mass)
atom_feat.append(vdw)
# atom_feat.append(pcharge)
atom_features.append(atom_feat)
for bond in mol.GetBonds():
g.add_edge(bond.GetBeginAtomIdx(), bond.GetEndAtomIdx())
g.ndata["feat"] = torch.FloatTensor(atom_features)
# Bond features
bond_features = []
for bond in mol.GetBonds():
bond_feat = []
bond_type = [0] * len(BOND_TYPES)
bond_type[BOND_TYPES.index(bond.GetBondType())] = 1
bond_stereo = [0] * len(BOND_STEREO)
bond_stereo[BOND_STEREO.index(bond.GetStereo())] = 1
bond_feat.extend(bond_type)
bond_feat.extend(bond_stereo)
bond_feat.append(float(bond.GetIsConjugated()))
bond_feat.append(float(bond.IsInRing()))
bond_features.append(bond_feat)
g.edata["feat"] = torch.FloatTensor(bond_features)
return g