예제 #1
0
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)
예제 #2
0
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)
예제 #3
0
    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")