def get_hybridization(self, atom: Atom):
hyb = str(atom.GetHybridization())
hyb = hyb.strip("SP")
hyb = "1" if hyb == "" else hyb
assert hyb in ["1", "2",
"3"], "No hybridization assigned for the atom!"
return hyb
def atom_to_feat(self, atm: Chem.Atom, owning_mol: Chem.Mol, idx):
# nb the func GetOwningMol could not be used for mol as would return different object each time
this_atms_idx = atm.GetIdx()
assert idx == this_atms_idx
feat = torch.zeros(len(self), dtype=torch.float32)
# One hot encoding of element
try:
feat[self.atms_to_idx[atm.GetSymbol()]] = 1.
except KeyError as ex:
warnings.warn(f"Ignoring the symbol {atm.GetSymbol()}")
idx_up_to = self.number_atom_options
# Atomic Number
feat[idx_up_to] = float(atm.GetAtomicNum())
idx_up_to += 1
# Acceptor/donor
acceptor_ids, donor_ids = self.get_acceptor_and_donor_ids(owning_mol)
# Acceptor
feat[idx_up_to] = float(this_atms_idx in acceptor_ids)
idx_up_to += 1
# Donor
feat[idx_up_to] = float(this_atms_idx in donor_ids)
idx_up_to += 1
# Hybridization
hyb_idx = self.hybridization(atm.GetHybridization())
if hyb_idx is not None:
feat[idx_up_to + hyb_idx] = 1.
idx_up_to += self.number_hyb_options
# Aromatic
feat[idx_up_to] = float(atm.GetIsAromatic())
idx_up_to += 1
# Number of Hydrogens
feat[idx_up_to] = float(atm.GetNumImplicitHs())
idx_up_to += 1
return feat
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 hybridization(atom: Atom) -> HybridizationType:
"""Hybridisation (HybridizationType).
"""
return atom.GetHybridization()