def atom_features(atom: Chem.Atom):
return np.array(
encoding_onehot_unk(atom.GetSymbol(), [
'C', 'N', 'O', 'S', 'F', 'Si', 'P', 'Cl', 'Br', 'Mg', 'Na', 'Ca',
'Fe', 'As', 'Al', 'I', 'B', 'V', 'K', 'Tl', 'Yb', 'Sb', 'Sn', 'Ag',
'Pd', 'Co', 'Se', 'Ti', 'Zn', 'H', 'Li', 'Ge', 'Cu', 'Au', 'Ni',
'Cd', 'In', 'Mn', 'Zr', 'Cr', 'Pt', 'Hg', 'Pb', 'Unknown'
]) + encoding_onehot(atom.GetDegree(), [0, 1, 2, 3, 4, 5]) +
encoding_onehot_unk(atom.GetTotalNumHs(), [0, 1, 2, 3, 4]) +
encoding_onehot_unk(atom.GetImplicitValence(), [0, 1, 2, 3, 4, 5]) +
[atom.GetIsAromatic()])
def atom_to_string(atom: Chem.Atom) -> str:
return '{}|{}|{}|{}|{}|{}|{}|{}|{}|{}|{}|{}|{}|{}'.format(
atom.GetChiralTag(),
atom.GetDegree(),
atom.GetExplicitValence(),
atom.GetFormalCharge(),
atom.GetHybridization(),
atom.GetIsAromatic(),
atom.GetMass(),
atom.GetNumExplicitHs(),
atom.GetNumImplicitHs(),
atom.GetNumRadicalElectrons(),
atom.GetTotalDegree(),
atom.GetTotalNumHs(),
atom.GetTotalValence(),
atom.IsInRing()
)
def atom_feature(atom: Atom) -> Tuple[torch.Tensor, int]:
'''
Extract atom feature vector.
Returns: (feature_vec, feature_dim)
Features:
# item # dim
atomic number 25
exp-valence 7
imp-valence 4
hybridization 7
Hs 5
degree 6
formal charge 6
in ring 1
aromatic 1
mass / 100 1
(sum) 63
'''
index = torch.tensor([
ATOM_MAP[atom.GetAtomicNum()],
atom.GetExplicitValence(),
atom.GetImplicitValence(),
HYBRID_MAP[atom.GetHybridization()],
atom.GetTotalNumHs(),
atom.GetDegree(),
atom.GetFormalCharge() + 2, # -2 ~ 3
int(atom.IsInRing()),
int(atom.GetIsAromatic())
]) + OFFSET
vec = torch.zeros(ATOM_FDIM)
vec[index] = 1.0
vec[-1] = atom.GetMass() / 100
return vec, ATOM_FDIM
def degree(atom: Atom) -> int:
"""Number of directly bonded neighbours (int).
"""
return atom.GetDegree()