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