def all_connected_selector(session, models, results):
"""select all atoms connected to the current selection"""
# TODO: right mouse mode for this
cur_sel = selected_atoms(session)
bond_sel = selected_bonds(session)
for bond in bond_sel:
cur_sel = cur_sel.merge(Atoms(bond.atoms))
atoms = Atoms()
for atom in cur_sel:
if atom in atoms:
continue
elif atom.structure not in models:
continue
connected_atoms = get_fragment(atom, max_len=len(atom.structure.atoms))
atoms = atoms.merge(connected_atoms)
results.add_atoms(atoms)
def substituent_selection(session, sub_name, models, results):
#TODO: optimize - or cheat and use cython or something
#TODO: make it so it doesn't select things with just an H bonded to them
# e.g. sel OH should not select water molecules
# probably do a get_all_connected for each fragment and
# check if all_connected.subtract(Atoms[atom]).subtract(frag) leave just an H
atoms = Atoms()
sub = Substituent(sub_name)
chix_sub = ResidueCollection(sub).get_chimera(session)
sub_elements = sorted(chix_sub.atoms.elements.names)
sub_ranks = canonical_rank(Atoms(chix_sub.atoms))
sorted_sub_atoms = [
x for _, x in sorted(zip(sub_ranks, chix_sub.atoms),
key=lambda pair: pair[0])
]
length = len(sub.atoms)
#rank_time = 0
#frag_time = 0
for model in models:
if isinstance(model, AtomicStructure):
for atom in model.atoms:
#session.logger.info("checking groups on %s" % atom.atomspec)
for bonded_atom in atom.neighbors:
if bonded_atom.element.name != sub.atoms[0].element:
continue
#session.logger.info("fragment for %s" % bonded_atom.atomspec)
#frag_start = perf_counter()
frag = get_fragment(bonded_atom, atom, length)
#frag_stop = perf_counter()
#frag_time += frag_stop - frag_start
if frag.intersects(atoms):
continue
frag = frag.subtract(Atoms([atom]))
if len(frag) != length:
continue
elements = sorted(frag.elements.names)
if sub_elements != elements:
#session.logger.info("wrong elements")
continue
#rank_start = perf_counter()
frag_ranks = canonical_rank(frag)
#rank_stop = perf_counter()
#rank_time += rank_stop - rank_start
#session.logger.warning(", ".join(sub_elements))
#session.logger.warning("%s;\n%s" % (", ".join(str(x) for x in sorted(frag_ranks)), ", ".join(str(x) for x in sorted(sub_ranks))))
sorted_frag_atoms = [
x for _, x in sorted(zip(frag_ranks, frag.instances()),
key=lambda pair: pair[0])
]
#session.logger.warning("%s;\n%s" % (", ".join(x.atomspec for x in sorted_frag_atoms), ", ".join(x.name for x in sorted_sub_atoms)))
for a, b in zip(sorted_frag_atoms, sorted_sub_atoms):
#session.logger.info("%s %s" % (a.atomspec, b.name))
if a.element.name != b.element.name:
#session.logger.info("different elements")
break
#session.logger.info("bonded: %s; other: %s" % (bonded_atom.atomspec, atom.atomspec))
if a is not bonded_atom and len(a.neighbors) != len(
b.neighbors):
#session.logger.info("different num neighbors")
#session.logger.info("%s and %s" % (a.atomspec, b.name))
#session.logger.info("%i vs %i" % (len(a.neighbors), len(b.neighbors)))
break
elif a is bonded_atom and (len(a.neighbors) -
1) != len(b.neighbors):
#session.logger.info("first atom, different num neighbors")
#session.logger.info("%s and %s" % (a.atomspec, b.name))
#session.logger.info("%i vs %i" % (len(a.neighbors) - 1, len(b.neighbors)))
break
failed = False
for i, j, k in zip(
sorted([
aa.element.name for aa in a.neighbors
if ((aa is not atom and a is bonded_atom)
or a is not bonded_atom)
]),
sorted([bb.element.name
for bb in b.neighbors]),
sorted([
aa for aa in a.neighbors
if ((aa is not atom and a is bonded_atom)
or a is not bonded_atom)
]),
):
if i != j:
#session.logger.info("failed %s %s, %s" % (i, j, k.atomspec))
failed = True
break
if failed:
break
else:
atoms = atoms.merge(frag)
#session.logger.info("spent %f time fragmenting" % frag_time)
#session.logger.info("spent %f time ranking atoms" % rank_time)
results.add_atoms(atoms)
def select_pick(self, pick, set_sel):
if set_sel:
self.session.selection.clear()
if pick:
atoms = Atoms()
bonds = Bonds()
if not hasattr(pick, "__iter__"):
pick = [pick]
first_selected = None
for p in pick:
if isinstance(p, PickedAtoms):
for atom in p.atoms:
if atom in atoms:
continue
connected_atoms = get_fragment(
atom, max_len=atom.structure.num_atoms)
atoms = atoms.merge(connected_atoms)
if first_selected is None:
first_selected = atoms[0].selected
elif isinstance(p, PickedAtom):
if p.atom in atoms:
continue
connected_atoms = get_fragment(
p.atom, max_len=p.atom.structure.num_atoms)
atoms = atoms.merge(connected_atoms)
if first_selected is None:
first_selected = atoms[0].selected
elif isinstance(p, PickedBonds):
for bond in p.bonds:
if bond.atoms[0] in atoms:
continue
connected_atoms = get_fragment(
bond.atoms[0], max_len=bond.structure.num_atoms)
atoms = atoms.merge(connected_atoms)
if first_selected is None:
first_selected = atoms[0].selected
elif isinstance(p, PickedBond):
if p.bond.atoms[0] in atoms:
continue
connected_atoms = get_fragment(
p.bond.atoms[0], max_len=p.bond.structure.num_atoms)
atoms = atoms.merge(connected_atoms)
if first_selected is None:
first_selected = atoms[0].selected
elif isinstance(p, PickedResidues):
for res in p.residues:
for atom in res:
if atom in atoms:
continue
connected_atoms = get_fragment(
atom, max_len=res.structure.num_atoms)
atoms = atoms.merge(connected_atoms)
if first_selected is None:
first_selected = atoms[0].selected
elif isinstance(p, PickedResidue):
for atom in p.residue.atoms:
if atoms in atoms:
continue
connected_atoms = get_fragment(
atom, max_len=p.residue.structure.num_atoms)
atoms = atoms.merge(connected_atoms)
if first_selected is None:
first_selected = atoms[0].selected
for atom in atoms:
for neighbor, bond in zip(atom.neighbors, atom.bonds):
bonds = bonds.merge(Bonds([bond]))
if first_selected is not None:
atoms.selected = not first_selected
bonds.selected = not first_selected
else:
run(self.session, "select clear")