"""Classes for handling arbitrary sets of atoms.""" import _modeller import sys from modeller import physical, group_restraints, energy_data, modfile from modeller.energy_profile import EnergyProfile from modeller.util import modutil from modeller.util.matrix import matrix_to_list __docformat__ = "epytext en" class SelectionIterator(object): """Utility class to iterate over all atoms in a selection.""" def __init__(self, mdl, seldict): self.mdl = mdl self.seliter = seldict.__iter__() def __iter__(self): return self if sys.version_info[0] >= 3: def __next__(self): obj = next(self.seliter) return self.mdl.atoms[obj - 1] else: def next(self): obj = self.seliter.next() return self.mdl.atoms[obj - 1] class selection(object): """An arbitrary set of atoms. 'atom' or 'residue' objects can be added to or removed from the selection, and selections can be manipulated in the same way as Python-standard sets - e.g. union, intersection, difference, or the equivalent |, & and - operators. Note that C{obj in sel} is only True if ALL atoms in 'obj' are in the selection sel. If you want to check for partial selection, use C{len(sel.intersection([obj])) > 0}.""" # Global reference to the DOPE restraints files, so that they're loaded # only once dope_restraints = None dopehr_restraints = None def __init__(self, *atoms): self.__selection = {} self.__mdl = None self.union_update(atoms) def get_atom_indices(self): """@return: the integer indices of all atoms in this selection""" if sys.version_info[:2] == (2,3): keys = self.__selection.keys() keys.sort() else: keys = sorted(self.__selection.keys()) return (keys, self.__mdl) def get_model(self): """@return: the model object which all selected atoms belong to. @rtype: model""" return self.__mdl def __repr__(self): suffix = "s" num = len(self.__selection) if num == 1: suffix = "" return "Selection of %d atom%s" % (num, suffix) def __str__(self): return "<%s>" % repr(self) def __len__(self): return len(self.__selection) def __check_object(self, obj): if not hasattr(obj, 'get_atom_indices'): raise TypeError(("Invalid object %s for selection - try atom, " + "residue, model, selection objects") % repr(obj)) def __contains__(self, obj): self.__check_object(obj) (inds, mdl) = obj.get_atom_indices() if mdl is not self.__mdl: return False for ind in inds: if ind not in self.__selection: return False return True def __iter__(self): return SelectionIterator(self.__mdl, self.__selection) def __typecheck(self, obj): if not isinstance(obj, selection): raise TypeError("Must use selection objects here") def __typecoerce(self, obj): if isinstance(obj, selection): return obj else: return selection(*obj) def issubset(self, t): """Tests whether every element is also in t""" t = self.__typecoerce(t) for obj in self: if obj not in t: return False return True def __le__(self, t): self.__typecheck(t) return self.issubset(t) def issuperset(self, t): """Tests whether every element in t is also in this set""" t = self.__typecoerce(t) for obj in t: if obj not in self: return False return True def __ge__(self, t): self.__typecheck(t) return self.issuperset(t) def __eq__(self, t): self.__typecheck(t) return self.issubset(t) and t.issubset(self) def __lt__(self, t): self.__typecheck(t) return self.issubset(t) and not t.issubset(self) def __gt__(self, t): self.__typecheck(t) return t.issubset(self) and not self.issubset(t) def add(self, obj): """Adds a new object (e.g. atom, residue) to the selection""" if not hasattr(obj, "get_atom_indices") \ and modutil.non_string_iterable(obj): for x in obj: self.add(x) else: self.__check_object(obj) (inds, mdl) = obj.get_atom_indices() if self.__mdl is None: self.__mdl = mdl elif self.__mdl is not mdl: raise ValueError("All atoms must be in the same model!") for ind in inds: self.__selection[ind] = None def remove(self, obj): """Removes an object (e.g. atom, residue) from the selection; raises KeyError if not present""" if not hasattr(obj, "get_atom_indices") \ and modutil.non_string_iterable(obj): for x in obj: self.remove(x) else: if obj not in self: raise KeyError(obj) (inds, mdl) = obj.get_atom_indices() if self.__mdl is not mdl: raise ValueError("All atoms must be in the same model!") for ind in inds: del self.__selection[ind] def discard(self, obj): """Removes an object (e.g. atom, residue) from the selection, if present""" try: self.remove(obj) except KeyError: pass def union(self, t): """@return: a new selection containing objects from both selections. @rtype: L{selection}""" newobj = selection() # Handle the case where both selections are empty newobj.__mdl = self.__mdl for obj in self: newobj.add(obj) for obj in t: newobj.add(obj) return newobj def __or__(self, t): self.__typecheck(t) return self.union(t) def intersection(self, t): """@return: a new selection containing objects common to both selections. @rtype: L{selection}""" t = self.__typecoerce(t) newobj = selection() # Handle the case where both selections are empty newobj.__mdl = self.__mdl for obj in self: if obj in t: newobj.add(obj) return newobj def __and__(self, t): self.__typecheck(t) return self.intersection(t) def difference(self, t): """@return: a new selection containing objects in this selection but not in t. @rtype: L{selection}""" t = self.__typecoerce(t) newobj = selection() # Handle the case where both selections are empty newobj.__mdl = self.__mdl for obj in self: if obj not in t: newobj.add(obj) return newobj def __sub__(self, t): self.__typecheck(t) return self.difference(t) def symmetric_difference(self, t): """Returns a new selection containing objects in either this selection or t, but not both""" t = self.__typecoerce(t) newobj = selection() # Handle the case where both selections are empty newobj.__mdl = self.__mdl for obj in self: if obj not in t: newobj.add(obj) for obj in t: if obj not in self: newobj.add(obj) return newobj def __xor__(self, t): self.__typecheck(t) return self.symmetric_difference(t) def copy(self): """@return: a copy of this selection. @rtype: L{selection}""" newobj = selection() newobj.__selection = self.__selection.copy() newobj.__mdl = self.__mdl return newobj def union_update(self, t): for obj in t: self.add(obj) return self def __ior__(self, t): self.__typecheck(t) return self.union_update(t) def intersection_update(self, t): t = self.__typecoerce(t) c = self.copy() for obj in c: if obj not in t: self.remove(obj) return self def __iand__(self, t): self.__typecheck(t) return self.intersection_update(t) def difference_update(self, t): t = self.__typecoerce(t) c = self.copy() for obj in c: if obj in t: self.remove(obj) return self def __isub__(self, t): self.__typecheck(t) return self.difference_update(t) def symmetric_difference_update(self, t): t = self.__typecoerce(t) c = self.copy() for obj in t: if obj in c: self.remove(obj) else: self.add(obj) return self def __ixor__(self, t): self.__typecheck(t) return self.symmetric_difference_update(t) def clear(self): """Remove all atoms from the selection""" self.__mdl = None self.__selection = {} def __require_indices(self): """Get atom indices, and fail if there are none""" (inds, mdl) = self.get_atom_indices() if mdl is None or len(inds) == 0: raise ValueError("Selection contains no atoms") return (inds, mdl) def translate(self, vector): """Translate the selection by the given vector""" (inds, mdl) = self.__require_indices() _modeller.mod_selection_translate(mdl.modpt, inds, vector) def rotate_origin(self, axis, angle): """Rotate the selection about the given axis through the origin, by the given angle (in degrees)""" (inds, mdl) = self.__require_indices() if sum([x * x for x in axis]) < 1e-8: raise ValueError("axis must be a vector of non-zero length") _modeller.mod_selection_rotate_axis(mdl.modpt, inds, axis, angle) def rotate_mass_center(self, axis, angle): """Rotate the selection about the given axis through the mass center, by the given angle (in degrees)""" com = self.mass_center self.translate([-a for a in com]) try: self.rotate_origin(axis, angle) finally: # If rotate_origin fails, at least restore the original conformation self.translate(com) def transform(self, matrix): """Transform the selection coordinates with the given matrix""" (inds, mdl) = self.__require_indices() lst = matrix_to_list(matrix) _modeller.mod_selection_transform(mdl.modpt, inds, lst) def find_atoms(self, restyp, atom_names, min_selected): if min_selected <= 0: raise ValueError("min_selected must be at least 1") mdl = self.get_model() if mdl is None or len(self) == 0: raise ValueError("Selection contains no atoms") from modeller import model_topology return model_topology.FindAtoms(mdl, restyp, atom_names, self.__check_selected, min_selected) def find_chi1_dihedrals(self, restyp, min_selected): return self.__find_dihedrals(restyp, 5, min_selected) def find_chi2_dihedrals(self, restyp, min_selected): return self.__find_dihedrals(restyp, 6, min_selected) def find_chi3_dihedrals(self, restyp, min_selected): return self.__find_dihedrals(restyp, 7, min_selected) def find_chi4_dihedrals(self, restyp, min_selected): return self.__find_dihedrals(restyp, 8, min_selected) def __find_dihedrals(self, restyp, dihedral_type, min_selected): if min_selected <= 0: raise ValueError("min_selected must be at least 1") mdl = self.get_model() if mdl is None or len(self) == 0: raise ValueError("Selection contains no atoms") from modeller import model_topology return model_topology.FindDihedrals(mdl, restyp, dihedral_type, self.__check_selected, min_selected) def __check_selected(self, atoms, min_selected): num_selected = 0 min_selected = min(min_selected, len(atoms)) for a in atoms: if a in self: num_selected += 1 if num_selected >= min_selected: return True return False def write(self, file, model_format='PDB', no_ter=False, extra_data=""): """Write selection coordinates to a file""" (inds, mdl) = self.__require_indices() fh = modfile._get_filehandle(file, 'w') return _modeller.mod_model_write(mdl.modpt, mdl.env.libs.modpt, inds, fh.file_pointer, model_format, no_ter, False, extra_data) def by_residue(self): """@return: a new selection, in which any residues in the existing selection, that have at least one selected atom, are now entirely selected. @rtype: L{selection}""" return self.extend_by_residue(0) def extend_by_residue(self, extension): """@return: a new selection, in which any residues with at least one selected atom in the existing selection are now entirely selected. Additionally, C{extension} residues around each selected residue are selected. @rtype: L{selection}""" newobj = self.copy() (inds, mdl) = self.get_atom_indices() if mdl is not None: for (n, res) in enumerate(mdl.residues): if len(self.intersection([res])) > 0: for i in range(max(0, n-extension), min(len(mdl.residues), n+extension+1)): newobj.add(mdl.residues[i]) return newobj def only_sidechain(self): """@return: a new selection, containing only sidechain atoms from the current selection. @rtype: L{selection}""" return self - self.only_mainchain() def only_mainchain(self): """@return: a new selection, containing only mainchain atoms from the current selection. @rtype: L{selection}""" return self.only_atom_types('O OT1 OT2 OXT C CA N') def only_atom_types(self, atom_types): """@return: a new selection, containing only atoms from the current selection of the given space-separated type(s) (e.g. 'CA CB'). @rtype: L{selection}""" if modutil.non_string_iterable(atom_types): atom_types = " ".join(atom_types) return self.__filter(_modeller.mod_selection_atom_types, atom_types) def only_residue_types(self, residue_types): """@return: a new selection, containing only atoms from the current selection in residues of the given space-separated type(s) (e.g. 'ALA ASP'). @rtype: L{selection}""" if modutil.non_string_iterable(residue_types): residue_types = " ".join(residue_types) libs = self.__get_libs() return self.__filter(_modeller.mod_selection_residue_types, residue_types, libs) def only_std_residues(self): """@return: a new selection, containing only atoms in standard residue types (i.e. everything but HETATM). @rtype: L{selection}""" libs = self.__get_libs() return self.__filter(_modeller.mod_selection_std_residues, libs) def only_no_topology(self): """@return: a new selection, containing only atoms in residues that have no defined topology. @rtype: L{selection}""" libs = self.__get_libs() return self.__filter(_modeller.mod_selection_no_topology, libs) def only_het_residues(self): """@return: a new selection, containing only atoms in HETATM residues. @rtype: L{selection}""" libs = self.__get_libs() return self.__filter(_modeller.mod_selection_het_residues, libs) def only_water_residues(self): """@return: a new selection, containing only atoms in water residues. @rtype: L{selection}""" libs = self.__get_libs() return self.__filter(_modeller.mod_selection_water_residues, libs) def only_defined(self): """@return: a new selection, containing only atoms with defined coordinates. @rtype: L{selection}""" return self.__filter(_modeller.mod_selection_defined) def select_sphere(self, radius): """@return: a new selection, containing all atoms within the given distance from any atom in the current selection. @rtype: L{selection}""" s = selection() for x in self: s.add(x.select_sphere(radius)) return s def mutate(self, residue_type): """Mutate selected residues""" (inds, mdl) = self.__require_indices() return _modeller.mod_selection_mutate(mdl.modpt, mdl.env.libs.modpt, inds, residue_type) def randomize_xyz(self, deviation): """Randomize coordinates""" (inds, mdl) = self.__require_indices() return _modeller.mod_randomize_xyz(mdl.modpt, mdl.env.libs.modpt, inds, deviation) def superpose(self, mdl2, aln, fit=True, superpose_refine=False, rms_cutoff=3.5, reference_atom='', reference_distance=3.5, refine_local=True, swap_atoms_in_res=''): """Superpose the input model on this selection, given an alignment of the models""" from modeller import superpose (inds, mdl) = self.__require_indices() retval = _modeller.mod_superpose(mdl.modpt, mdl2.modpt, aln.modpt, mdl.env.libs.modpt, inds, swap_atoms_in_res, reference_atom, reference_distance, superpose_refine, fit, refine_local, rms_cutoff) return superpose.SuperposeData(*retval) def rotate_dihedrals(self, deviation, change, dihedrals=('PHI', 'PSI', 'CHI1', 'CHI2', 'CHI3', 'CHI4')): """Optimize or randomize dihedral angles""" (inds, mdl) = self.__require_indices() return _modeller.mod_rotate_dihedrals(mdl.modpt, mdl.env.libs.modpt, inds, deviation, change, dihedrals) def hot_atoms(self, pick_hot_cutoff, residue_span_range=(0, 99999), viol_report_cut=physical.values(default=4.5, chi1_dihedral=999, chi2_dihedral=999, chi3_dihedral=999, chi4_dihedral=999, chi5_dihedral=999, phi_psi_dihedral=6.5, nonbond_spline=999, accessibility=999, density=999, gbsa=999, em_density=999), schedule_scale=None, edat=None): """@return: a new selection containing all atoms violating restraints. @rtype: L{selection}""" (inds, mdl) = self.__require_indices() if edat is None: edat = mdl.env.edat if schedule_scale is None: schedule_scale = mdl.env.schedule_scale func = _modeller.mod_selection_hot_atoms newinds = func(mdl.modpt, edat.modpt, mdl.env.libs.modpt, inds, residue_span_range, pick_hot_cutoff, viol_report_cut, schedule_scale) newobj = selection() newobj.__mdl = mdl newobj.__selection = dict.fromkeys(newinds) return newobj def objfunc(self, edat=None, residue_span_range=(0,99999), schedule_scale=physical.values(default=1.0)): """Get just the objective function value, without derivatives""" (inds, mdl) = self.__require_indices() if edat is None: edat = mdl.env.edat return _modeller.mod_selection_objfunc(mdl.modpt, edat.modpt, mdl.env.libs.modpt, inds, residue_span_range, schedule_scale) def energy(self, asgl_output=False, normalize_profile=False, residue_span_range=(0, 99999), output='LONG', file='default', viol_report_cut=physical.values(default=4.5, chi1_dihedral=999, chi2_dihedral=999, chi3_dihedral=999, chi4_dihedral=999, chi5_dihedral=999, phi_psi_dihedral=6.5, nonbond_spline=999, accessibility=999, density=999, gbsa=999, em_density=999), viol_report_cut2=physical.values(default=2.0), smoothing_window=3, schedule_scale=None, edat=None): """Evaluate the objective function given restraints""" (inds, mdl) = self.__require_indices() if edat is None: edat = mdl.env.edat if schedule_scale is None: schedule_scale = mdl.env.schedule_scale func = _modeller.mod_energy (molpdf, terms) = func(mdl.modpt, edat.modpt, mdl.env.libs.modpt, inds, asgl_output, normalize_profile, residue_span_range, output, file, smoothing_window, viol_report_cut, viol_report_cut2, schedule_scale) terms = physical.from_list(terms) return (molpdf, terms) def get_dope_potential(self): """Get the spline data for the DOPE statistical potential""" (inds, mdl) = self.__require_indices() if not selection.dope_restraints: selection.dope_restraints = \ group_restraints(mdl.env, classes='${LIB}/atmcls-mf.lib', parameters='${LIB}/dist-mf.lib') return selection.dope_restraints def get_dopehr_potential(self): """Get the spline data for the DOPE-HR statistical potential""" (inds, mdl) = self.__require_indices() if not selection.dopehr_restraints: selection.dopehr_restraints = \ group_restraints(mdl.env, classes='${LIB}/atmcls-mf.lib', parameters='${LIB}/dist-mfhr.lib') return selection.dopehr_restraints def get_dope_energy_data(self): """Get ideal energy_data terms for DOPE potential evaluations""" return energy_data(contact_shell=15.0, dynamic_modeller=True, dynamic_lennard=False, dynamic_sphere=False, excl_local=(False, False, False, False)) def assess_dope(self, **vars): """Assess the selection with the DOPE potential""" return self._dope_energy(self.get_dope_potential(), "DOPE", **vars) def assess_dopehr(self, **vars): """Assess the selection with the DOPE-HR potential""" return self._dope_energy(self.get_dopehr_potential(), "DOPE-HR", **vars) def get_energy_profile(self, physical_type, edat=None): """Get a per-residue energy profile, plus the number of restraints on each residue, and the RMS minimum and heavy violations""" (inds, mdl) = self.__require_indices() if edat is None: edat = mdl.env.edat scaln = physical.values(default=0.) scaln[physical_type] = 1. prof = _modeller.mod_rms_profile(mdl.modpt, edat.modpt, mdl.env.libs.modpt, inds, (1, 9999), True, False, physical_type.get_type(), scaln) return EnergyProfile(*prof) def get_dope_profile(self): """Get a per-residue DOPE energy profile""" return self._dope_profile(self.get_dope_potential()) def get_dopehr_profile(self): """Get a per-residue DOPE-HR energy profile""" return self._dope_profile(self.get_dopehr_potential()) def _dope_profile(self, gprsr): """Internal function to get a DOPE or DOPE-HR energy profile""" mdl = self.__mdl edat = self.get_dope_energy_data() old_gprsr = mdl.group_restraints mdl.group_restraints = gprsr try: prof = self.get_energy_profile(physical.nonbond_spline, edat) finally: mdl.group_restraints = old_gprsr return prof def assess(self, assessor, output='SHORT NO_REPORT', **vars): """Assess with the given assessor object (e.g. soap_loop.Scorer())""" print(">> Model assessment by %s" % assessor.name) molpdf, terms = assessor._assess(self, output=output, **vars) print("%s : %12.6f" % (assessor.name, molpdf)) return molpdf def _dope_energy(self, gprsr, name, output='SHORT NO_REPORT', residue_span_range=(1, 9999), schedule_scale=physical.values(default=0., nonbond_spline=1.), **vars): """Internal function to do DOPE or DOPE-HR assessment""" mdl = self.__mdl print(">> Model assessment by %s potential" % name) edat = self.get_dope_energy_data() old_gprsr = mdl.group_restraints mdl.group_restraints = gprsr try: (molpdf, terms) = \ self.energy(edat=edat, residue_span_range=residue_span_range, output=output, schedule_scale=schedule_scale, **vars) finally: mdl.group_restraints = old_gprsr print("%s score : %12.6f" % (name, molpdf)) return molpdf def debug_function(self, residue_span_range=(0, 99999), debug_function_cutoff=(0.01, 0.001, 0.1), detailed_debugging=False, schedule_scale=None, edat=None): """Test code self-consistency""" (inds, mdl) = self.__require_indices() if edat is None: edat = mdl.env.edat if schedule_scale is None: schedule_scale = mdl.env.schedule_scale func = _modeller.mod_debug_function return func(mdl.modpt, edat.modpt, mdl.env.libs.modpt, inds, residue_span_range, debug_function_cutoff, detailed_debugging, schedule_scale) def unbuild(self): """Undefine all coordinates""" (inds, mdl) = self.__require_indices() return _modeller.mod_selection_unbuild(mdl=mdl.modpt, sel1=inds) def __get_libs(self): if self.__mdl: return self.__mdl.env.libs.modpt else: return None def __filter(self, func, *args): newobj = selection() (inds, mdl) = self.get_atom_indices() if mdl is not None: newinds = func(mdl.modpt, inds, *args) newobj.__mdl = mdl newobj.__selection = dict.fromkeys(newinds) return newobj def __get_com(self): (inds, mdl) = self.__require_indices() return _modeller.mod_selection_com_get(mdl.modpt, inds) def __set_com(self, val): com = self.mass_center self.translate([val[0] - com[0], val[1] - com[1], val[2] - com[2]]) def __get_x(self): return self.__get_com()[0] def __get_y(self): return self.__get_com()[1] def __get_z(self): return self.__get_com()[2] def __set_x(self, val): self.translate([val - self.x, 0, 0]) def __set_y(self, val): self.translate([0, val - self.y, 0]) def __set_z(self, val): self.translate([0, 0, val - self.z]) mass_center = property(__get_com, __set_com, doc="Coordinates of mass center") x = property(__get_x, __set_x, doc="x coordinate of mass center") y = property(__get_y, __set_y, doc="y coordinate of mass center") z = property(__get_z, __set_z, doc="z coordinate of mass center")