def get_atom_type(self, atom: Chem.Atom) -> int:
"""
Get the atom type (represented as the index in self.atom_types) of the
given atom `atom`
Args:
atom (Chem.Atom):
The input atom
Returns:
int:
The atom type as int
"""
atom_symbol = atom.GetSymbol()
atom_charge = atom.GetFormalCharge()
atom_hs = atom.GetNumExplicitHs()
return self.atom_types.index((atom_symbol, atom_charge, atom_hs))
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()
)