def set_perpendicular(self, *args): atoms = selected_atoms(self.session) if len(atoms) == 0: self.session.logger.error("no atoms selected") return self.perpendiculars = {} self.perp_centers = {} models = set(atom.structure for atom in atoms) for model in models: atom_coords = [] for atom in atoms: if atom.structure is model: atom_coords.append(atom.coord) if len(atom_coords) < 3: self.session.logger.error("fewer than 3 atoms selected on %s" % model.atomspec) continue xyz = np.array(atom_coords) xyz -= np.mean(atom_coords, axis=0) R = np.dot(xyz.T, xyz) u, s, vh = np.linalg.svd(R, compute_uv=True) vector = u[:, -1] self.perpendiculars[model] = vector self.perp_centers[model] = np.mean(atom_coords, axis=0)
def run_script(session): from chimerax.atomic import selected_atoms from chimerax.build_structure import modify_atom from chimerax.build_structure.mod import ParamError from chimerax.core.errors import UserError sel = selected_atoms(session) if len(sel) != 1: raise UserError('Please select a single atom!') sel = sel[0] current_num_bonds = len(sel.neighbors) current_color = sel.color try: modify_atom(sel, sel.element, current_num_bonds + 1, connect_back=False, res_name=sel.residue.name) sel.color = current_color except ParamError as e: # If modify_atom throws an error at the previous step, it will have deleted # any attached hydrogens and not put them back. We need to put them back here. modify_atom(sel, sel.element, current_num_bonds, connect_back=False, res_name=sel.residue.name) sel.color = current_color raise UserError(str(e))
def run_bond(self, *args): # TODO: switch to `bond sel` in 1.2 sel = selected_atoms(self.session) halfbond = self.bond_halfbond.checkState() == Qt.Checked self.settings.bond_halfbond = halfbond if not halfbond: color = self.bond_color.get_color() color = tuple(x / 255. for x in color) self.settings.bond_color = color radius = self.bond_radius.value() self.settings.bond_radius = radius for b in selected_bonds(self.session): b.halfbond = halfbond if not halfbond: b.color = np.array([int(x * 255) for x in color]) b.radius = radius for i, a1 in enumerate(sel): for a2 in sel[:i]: if a1.structure is a2.structure and a2 not in a1.neighbors: new_bond = a1.structure.new_bond(a1, a2) new_bond.halfbond = halfbond if not halfbond: new_bond.color = np.array( [int(x * 255) for x in color]) new_bond.radius = radius
def selected_vertices(session): from chimerax.atomic import selected_atoms atoms = selected_atoms(session) from numpy import array, bool catoms = atoms.filter(array([hasattr(a, 'polygon') for a in atoms], bool)) return catoms
def callback(self, session): from chimerax.atomic import selected_atoms a1, a2 = selected_atoms(session) command = "dist %s %s" % (a1.string(style="command line"), a2.string(style="command line")) from chimerax.core.commands import run run(session, command)
def set_ligand(self, *args): self.ligands = {} for atom in selected_atoms(self.session): if atom.structure not in self.ligands: self.ligands[atom.structure] = [] self.ligands[atom.structure].append(atom) self.session.logger.status("set ligand to current selection")
def link_consecutive(self, event): s = self.session from chimerax.atomic import selected_atoms atoms1 = selected_atoms(s) a2 = self.picked_marker(event, select=True) if a2 is None or len(atoms1) != 1: return False a1 = atoms1[0] if a1.structure != a2.structure: s.logger.status('Cannot connect atoms from different molecules') return False if a1 is a2 or a1.connects_to(a2): return False ms = _mouse_marker_settings(self.session) from .markers import create_link b = create_link(a1, a2, radius=ms['link radius'], rgba=ms['link color'], log=True) s.logger.status('Made connection, distance %.3g' % b.length) return True
def _mouse_place_marker(session, center, link_to_selected=False, select=True, log=True): m = _mouse_markerset(session) ms = _mouse_marker_settings(session) a = m.create_marker(center, ms['marker color'], ms['marker radius'], ms['next_marker_num']) if log: _log_place_marker(m, center, ms['marker color'], ms['marker radius']) ms['next_marker_num'] += 1 session.logger.status('Placed marker') if link_to_selected: from chimerax.atomic import selected_atoms atoms = selected_atoms(session) if len(atoms) == 1: al = atoms[0] if a.structure == al.structure and a is not al: from .markers import create_link create_link(al, a, radius=ms['link radius'], rgba=ms['link color'], log=log) if select: session.selection.clear() a.selected = True
def mouse_down(self, event): MouseMode.mouse_down(self, event) if self.action(event) == 'rotate': self._set_z_rotation(event) if self.move_atoms: from chimerax.atomic import selected_atoms self._atoms = selected_atoms(self.session) self._undo_start()
def _create_distance(self): from chimerax.atomic import selected_atoms sel_atoms = selected_atoms(self.session) if len(sel_atoms) != 2: from chimerax.core.errors import UserError raise UserError("Exactly two atoms must be selected!") from chimerax.core.commands import run run(self.session, "distance %s %s" % tuple(a.string(style="command") for a in sel_atoms))
def wheel(self, event): d = event.wheel_value() if self.move_atoms: from chimerax.atomic import selected_atoms self._atoms = selected_atoms(self.session) if self.action(event) == 'rotate': self._rotate((0, 1, 0), 10 * d) else: self._translate((0, 0, 100 * d))
def update_key_atoms(self): selection = selected_atoms(self.session) if not selection.single_structure: raise RuntimeError("selected atoms must be on the same model") else: self.key_atomspec = selection self.tool_window.status("key atoms set to %s" % " ".join(atom.atomspec for atom in selection))
def incr_b_factor(session, b_add, atoms=None): B_MAX = 500 if atoms is None: from chimerax.atomic import selected_atoms atoms = selected_atoms(session) if any(atoms.bfactors + b_add < 0): raise UserError( 'Applying this command would reduce the B-factor of at least one atom to below zero.' ) atoms.bfactors += b_add atoms[atoms.bfactors > B_MAX].bfactors = B_MAX
def _ms_sel_changed(self, *args): from chimerax.atomic import selected_atoms sel_atoms = selected_atoms(self.session) if len(sel_atoms) != 1: return a = sel_atoms[0] self._ms_update_atom_name(a) from .mod import unknown_res_name res_name = unknown_res_name(a.residue) self.ms_mod_edit.setText(res_name) self.ms_res_new_name.setText(res_name)
def run_script(session): from chimerax.atomic import selected_atoms from chimerax.core.errors import UserError sel = selected_atoms(session) if len(sel) != 2: raise UserError('Must have exactly two atoms selected!') us = sel.unique_structures if len(us) != 1: raise UserError('Both atoms must be from the same structure!') m = us[0] from chimerax.atomic.struct_edit import add_bond add_bond(*sel)
def _fit_atoms(self): m = self._object_menu.value if m == 'selected atoms': from chimerax.atomic import selected_atoms return selected_atoms(self.session) from chimerax.atomic import Structure if isinstance(m, Structure): return m.atoms return None
def selection_changed(self, *args): selection = selected_atoms(self.session) new_atoms = [atom for atom in selection if atom not in self._selection] if len(new_atoms) > 1: self._selection = [] else: for atom in self._selection: if atom not in selection: self._selection.remove(atom) self._selection.extend(new_atoms)
def change_bond_length(self, *args): dist = self.bond_distance.value() atom_pairs = [] sel = selected_atoms(self.session) if len(sel) == 2 and sel[0].structure is sel[1].structure: atom_pairs.append(sel) for bond in selected_bonds(self.session): if not all(atom in sel for atom in bond.atoms): atom_pairs.append(bond.atoms) for bond in selected_pseudobonds(self.session): if not all(atom in sel for atom in bond.atoms): atom_pairs.append(bond.atoms) for pair in atom_pairs: atom1, atom2 = pair frag1 = get_fragment(atom1, stop=atom2, max_len=atom1.structure.num_atoms) frag2 = get_fragment(atom2, stop=atom1, max_len=atom1.structure.num_atoms) v = atom2.coord - atom1.coord cur_dist = np.linalg.norm(v) change = dist - cur_dist if self.move_fragment.currentText() == "both": change = 0.5 * change frag1.coords -= change * v / cur_dist frag2.coords += change * v / cur_dist elif self.move_fragment.currentText() == "smaller": if len(frag1) < len(frag2) or (len(frag1) == len(frag2) and sum(frag1.elements.masses) < sum(frag2.elements.masses)): frag1.coords -= change * v / cur_dist else: frag2.coords += change * v / cur_dist elif self.move_fragment.currentText() == "larger": if len(frag1) > len(frag2) or (len(frag1) == len(frag2) and sum(frag1.elements.masses) > sum(frag2.elements.masses)): frag1.coords -= change * v / cur_dist else: frag2.coords += change * v / cur_dist
def _ms_update_atom_name(self, a=None): if a is None: from chimerax.atomic import selected_atoms sel_atoms = selected_atoms(self.session) if len(sel_atoms) != 1: return a = sel_atoms[0] new_element = self.ms_elements_button.text() from .mod import default_changed_name new_name = default_changed_name(a, new_element) self.ms_atom_name.setText(new_name) if new_name == a.name: self.ms_retain_atom_name.setChecked(True) else: self.ms_change_atom_name.setChecked(True)
def manual_cor(self, *args): selection = selected_atoms(self.session) models = {} for atom in selection: if atom.structure not in models: models[atom.structure] = [atom] else: models[atom.structure].append(atom) self.manual_center = {} for model in models: atoms = models[model] coords = np.array([atom.coord for atom in atoms]) self.manual_center[model] = np.mean(coords, axis=0)
def set_group(self, *args): atoms = selected_atoms(self.session) if len(atoms) == 0: self.session.logger.error("no atoms selected") return self.groups = {} models = set(atom.structure for atom in atoms) for model in models: atom_coords = [] for atom in atoms: if atom.structure is model: atom_coords.append(atom.coord) self.groups[model] = np.mean(atom_coords, axis=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)
def _ms_apply_cb(self): from chimerax.atomic import selected_atoms sel_atoms = selected_atoms(self.session) num_selected = len(sel_atoms) if num_selected != 1: raise UserError("You must select exactly one atom to modify.") a = sel_atoms[0] element_name = self.ms_elements_button.text() num_bonds = self.ms_bonds_button.text() cmd = "build modify %s %s %s" % (a.atomspec, element_name, num_bonds) geometry = self.ms_geom_button.text() if geometry != "N/A": cmd += " geometry " + geometry if not self.ms_retain_atom_name.isChecked(): new_name = self.ms_atom_name.text().strip() if not new_name: raise UserError("Must provide a name for the modified atom") if new_name != a.name: cmd += " name " + new_name if not self.ms_connect_back.isChecked(): cmd += " connectBack false" if not self.ms_element_color.isChecked(): cmd += " colorByElement false" if self.ms_res_mod.isChecked(): res_name = self.ms_mod_edit.text().strip() if not res_name: raise UserError("Must provided modified residue name") if res_name != a.residue.name: cmd += " resName " + res_name elif self.ms_res_new.isChecked(): res_name = self.ms_res_new_name.text().strip() if not res_name: raise UserError("Must provided new residue name") cmd += " resNewOnly true resName " + res_name run(self.session, cmd) if self.ms_focus.isChecked(): run(self.session, "view " + a.residue.atomspec)
def copy_atom_style_from(session, atoms, ref_residue): from chimerax.atomic import selected_atoms, Atom from chimerax.core.commands import run r = ref_residue atoms.draw_modes = r.atoms[0].draw_mode residues = atoms.unique_residues residues.ribbon_displays = r.ribbon_display residues.ribbon_hide_backbones = r.ribbon_hide_backbone current_sel = selected_atoms(session) session.selection.clear() atoms.selected = True ref_c = r.atoms[r.atoms.element_names == 'C'] if len(ref_c): atoms[atoms.element_names == 'C'].colors = ref_c[0].color else: run(session, "color sel bychain", log=False) run(session, "color sel byhetero", log=False) session.selection.clear() current_sel.selected = True
def intersect_selection(objects, session, undo_state, full_residues = False): atoms, bonds, pbonds, models = _atoms_bonds_models(objects, full_residues = full_residues) from chimerax import atomic selatoms = atomic.selected_atoms(session) subatoms = selatoms - atoms selbonds = atomic.selected_bonds(session) subbonds = selbonds - bonds selpbonds = atomic.selected_pseudobonds(session) subpbonds = selpbonds - pbonds from chimerax.atomic import Structure, PseudobondGroup selmodels = set(m for m in session.selection.models() if not isinstance(m, (Structure, PseudobondGroup))) submodels = selmodels.difference(models) undo_state.add(subatoms, "selected", subatoms.selected, False) undo_state.add(subbonds, "selected", subbonds.selected, False) undo_state.add(subpbonds, "selected", subpbonds.selected, False) subatoms.selected = False subbonds.selected = False subpbonds.selected = False for m in submodels: undo_state.add(m, "selected", m.selected, False) m.selected = False
from chimerax.core.commands import run run(session, "open 3fx2 ; sel ligand & aromatic") from chimerax.atomic import selected_atoms num_selected = len(selected_atoms(session)) if num_selected != 10: raise SystemExit("Selecting 3fx2 :FMN aromatic rings selected %d atoms instead of 10!" % num_selected) run(session, "sel phosphate") num_selected = len(selected_atoms(session)) if num_selected != 5: raise SystemExit("Selecting 3fx2 phoshates selected %d atoms instead of 5!" % num_selected)
def calc_cone(self, *args): self.settings.cone_option = self.cone_option.currentText() self.settings.radii = self.radii_option.currentText() self.settings.display_radii = self.display_radii.checkState( ) == Qt.Checked self.settings.display_cone = self.display_cone.checkState( ) == Qt.Checked if self.cone_option.currentText() == "Tolman (Unsymmetrical)": method = "tolman" else: method = self.cone_option.currentText() radii = self.radii_option.currentText() return_cones = self.display_cone.checkState() == Qt.Checked display_radii = self.display_radii.checkState() == Qt.Checked # self.table.setRowCount(0) for center_atom in selected_atoms(self.session): rescol = ResidueCollection(center_atom.structure) at_center = rescol.find_exact(AtomSpec(center_atom.atomspec))[0] if center_atom.structure in self.ligands: comp = Component( rescol.find([ AtomSpec(atom.atomspec) for atom in self.ligands[center_atom.structure] ]), to_center=rescol.find_exact(AtomSpec( center_atom.atomspec)), key_atoms=rescol.find(BondedTo(at_center)), ) else: comp = Component( rescol.find(NotAny(at_center)), to_center=rescol.find_exact(AtomSpec( center_atom.atomspec)), key_atoms=rescol.find(BondedTo(at_center)), ) cone_angle = comp.cone_angle( center=rescol.find(AtomSpec(center_atom.atomspec)), method=method, radii=radii, return_cones=return_cones, ) if return_cones: cone_angle, cones = cone_angle s = ".transparency 0.5\n" for cone in cones: apex, base, radius = cone s += ".cone %6.3f %6.3f %6.3f %6.3f %6.3f %6.3f %.3f open\n" % ( *apex, *base, radius) stream = BytesIO(bytes(s, "utf-8")) bild_obj, status = read_bild(self.session, stream, "Cone angle %s" % center_atom) self.session.models.add(bild_obj, parent=center_atom.structure) if display_radii: s = ".note radii\n" s += ".transparency 75\n" color = None for atom in comp.atoms: chix_atom = atom.chix_atom if radii.lower() == "umn": r = VDW_RADII[chix_atom.element.name] elif radii.lower() == "bondi": r = BONDI_RADII[chix_atom.element.name] if color is None or chix_atom.color != color: color = chix_atom.color rgb = [x / 255. for x in chix_atom.color] rgb.pop(-1) s += ".color %f %f %f\n" % tuple(rgb) s += ".sphere %f %f %f %f\n" % (*chix_atom.coord, r) stream = BytesIO(bytes(s, "utf-8")) bild_obj, status = read_bild(self.session, stream, "Cone angle radii") self.session.models.add(bild_obj, parent=center_atom.structure) row = self.table.rowCount() self.table.insertRow(row) name = QTableWidgetItem() name.setData(Qt.DisplayRole, center_atom.structure.name) self.table.setItem(row, 0, name) center = QTableWidgetItem() center.setData(Qt.DisplayRole, center_atom.atomspec) self.table.setItem(row, 1, center) ca = QTableWidgetItem() ca.setData(Qt.DisplayRole, "%.2f" % cone_angle) self.table.setItem(row, 2, ca) self.table.resizeColumnToContents(0) self.table.resizeColumnToContents(1) self.table.resizeColumnToContents(2)
def _sel_residues(session, models, results): from chimerax.atomic import selected_atoms results.add_atoms(selected_atoms(session).residues.unique().atoms)
def selected_markers(session): from chimerax import atomic atoms = atomic.selected_atoms(session) mask = [isinstance(a.structure, MarkerSet) for a in atoms] return atoms.filter(mask)
def show_rot_vec(self, *args): for model in self.session.models.list(type=Generic3DModel): if model.name == "rotation vector": model.delete() if self.display_rot_vec.checkState() == Qt.Unchecked: return selection = selected_atoms(self.session) if len(selection) == 0: return models = {} for atom in selection: if atom.structure not in models: models[atom.structure] = [atom] else: models[atom.structure].append(atom) if len(models.keys()) == 0: return if self.vector_option.currentText() == "axis": if self.axis.currentText() == "z": vector = np.array([0., 0., 1.]) elif self.axis.currentText() == "y": vector = np.array([0., 1., 0.]) elif self.axis.currentText() == "x": vector = np.array([1., 0., 0.]) elif self.vector_option.currentText() == "view axis": if self.view_axis.currentText() == "z": vector = self.session.view.camera.get_position().axes()[2] elif self.view_axis.currentText() == "y": vector = self.session.view.camera.get_position().axes()[1] elif self.view_axis.currentText() == "x": vector = self.session.view.camera.get_position().axes()[0] elif self.vector_option.currentText() == "bond": vector = self.bonds elif self.vector_option.currentText() == "perpendicular to plane": vector = self.perpendiculars elif self.vector_option.currentText() == "centroid of atoms": vector = self.groups elif self.vector_option.currentText() == "custom": x = self.vector_x.value() y = self.vector_y.value() z = self.vector_z.value() vector = np.array([x, y, z]) angle = np.deg2rad(self.angle.value()) center = {} for model in models: atoms = models[model] coords = np.array([atom.coord for atom in atoms]) center[model] = np.mean(coords, axis=0) if self.cor_button.currentText() == "automatic": if self.vector_option.currentText() == "perpendicular to plane": center = self.perp_centers elif self.vector_option.currentText() == "bond": center = self.bond_centers elif self.cor_button.currentText() == "select atoms": center = self.manual_center else: center = self.session.main_view.center_of_rotation for model in models: if isinstance(vector, dict): if model not in vector.keys(): continue else: v = vector[model] else: v = vector if isinstance(center, dict): if model not in center.keys(): continue else: c = center[model] else: c = center if self.vector_option.currentText( ) == "centroid of atoms" and self.cor_button.currentText( ) != "automatic": v = v - c if np.linalg.norm(v) == 0: continue residues = [] for atom in models[model]: if atom.residue not in residues: residues.append(atom.residue) v_c = c + v s = ".color red\n" s += ".arrow %10.6f %10.6f %10.6f %10.6f %10.6f %10.6f 0.2 0.4 0.7\n" % ( *c, *v_c) stream = BytesIO(bytes(s, 'utf-8')) bild_obj, status = read_bild(self.session, stream, "rotation vector") self.session.models.add(bild_obj, parent=model)