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