def index_to_bond(self, mol: Chem.Mol, begin_id: int, end_id: int,
idx: int):
""" Create a new bond in the given molecule `mol`
Args:
mol (Chem.Mol):
The molecule where the new bond should be created
begin_id, end_id (int):
The (starting and terminating) location of the new bond
idx (int):
The type of the new bond
"""
mol.AddBond(begin_id, end_id, self.bond_orders[idx])
def collapse_ring(self, mol: Chem.Mol) -> Chem.Mol:
"""
Collapses a ring(s) into a single dummy atom(s).
Stores data as JSON in the atom.
:param mol:
:return:
"""
self.store_positions(mol)
mol = Chem.RWMol(mol)
conf = mol.GetConformer()
center_idxs = []
morituri = []
old2center = defaultdict(list)
for atomset in mol.GetRingInfo().AtomRings():
morituri.extend(atomset)
neighs = []
neighbonds = []
bonds = []
xs = []
ys = []
zs = []
elements = []
# add elemental ring
c = mol.AddAtom(Chem.Atom('C'))
center_idxs.append(c)
central = mol.GetAtomWithIdx(c)
name = mol.GetProp('_Name') if mol.HasProp('_Name') else '???'
central.SetProp('_ori_name', name),
# get data for storage
for i in atomset:
old2center[i].append(c)
atom = mol.GetAtomWithIdx(i)
neigh_i = [a.GetIdx() for a in atom.GetNeighbors()]
neighs.append(neigh_i)
bond = [mol.GetBondBetweenAtoms(i, j).GetBondType().name for j in neigh_i]
bonds.append(bond)
pos = conf.GetAtomPosition(i)
xs.append(pos.x)
ys.append(pos.y)
zs.append(pos.z)
elements.append(atom.GetSymbol())
# store data in elemental ring
central.SetIntProp('_ori_i', -1)
central.SetProp('_ori_is', json.dumps(atomset))
central.SetProp('_neighbors', json.dumps(neighs))
central.SetProp('_xs', json.dumps(xs))
central.SetProp('_ys', json.dumps(ys))
central.SetProp('_zs', json.dumps(zs))
central.SetProp('_elements', json.dumps(elements))
central.SetProp('_bonds', json.dumps(bonds))
conf.SetAtomPosition(c, Point3D(*[sum(axis) / len(axis) for axis in (xs, ys, zs)]))
for atomset, center_i in zip(mol.GetRingInfo().AtomRings(), center_idxs):
# bond to elemental ring
central = mol.GetAtomWithIdx(center_i)
neighss = json.loads(central.GetProp('_neighbors'))
bondss = json.loads(central.GetProp('_bonds'))
for neighs, bonds in zip(neighss, bondss):
for neigh, bond in zip(neighs, bonds):
if neigh not in atomset:
bt = getattr(Chem.BondType, bond)
if neigh not in morituri:
mol.AddBond(center_i, neigh, bt)
else:
for other_center_i in old2center[neigh]:
if center_i != other_center_i:
if not mol.GetBondBetweenAtoms(center_i, other_center_i):
mol.AddBond(center_i, other_center_i, bt)
break
else:
raise ValueError(f'Cannot find what {neigh} became')
for i in sorted(set(morituri), reverse=True):
mol.RemoveAtom(self._get_new_index(mol, i))
return mol.GetMol()
