def atom_features(
atom: Chem.rdchem.Atom,
functional_groups: List[int] = None) -> List[Union[bool, int, float]]:
"""
Builds a feature vector for an atom.
:param atom: An RDKit atom.
:param functional_groups: A k-hot vector indicating the functional groups the atom belongs to.
:return: A list containing the atom features.
"""
features = onek_encoding_unk(atom.GetAtomicNum() - 1, ATOM_FEATURES['atomic_num']) + \
onek_encoding_unk(atom.GetTotalDegree(), ATOM_FEATURES['degree']) + \
onek_encoding_unk(int(atom.GetTotalNumHs()), ATOM_FEATURES['num_Hs']) + \
[atom.GetMass() * 0.01] # scaled to about the same range as other features
features += [
atom.IsInRingSize(3),
atom.IsInRingSize(4),
atom.IsInRingSize(5),
atom.IsInRingSize(6),
atom.IsInRingSize(7),
atom.IsInRingSize(8),
atom.IsInRingSize(9),
atom.IsInRingSize(10),
]
if functional_groups is not None:
features += functional_groups
return features
def atom_in_member_rings(atom: Chem.rdchem.Atom):
"""
Show each atom of the molecule is involved in 3, 4, 5, 6, 7, 8 member rings
"""
vector = [
0
] * 7 # If value is not in the [3, 4, 5, 6, 7, 8], then the final element in the vector is 1.
for index, member_ring in enumerate(range(3, 9)):
if atom.IsInRingSize(member_ring):
vector[index] = 1
if vector == [0] * 7:
vector[-1] = 1
return vector
