IMP logo
IMP Reference Guide  2.20.2
The Integrative Modeling Platform
system_tools.py
1 from __future__ import print_function, division
2 import IMP
3 import IMP.atom
5 import IMP.pmi
6 import IMP.pmi.tools
7 from collections import defaultdict
8 from math import pi
9 import os
10 import warnings
11 
12 # IMP doesn't statically define an atom type for CA atoms in modified residues
13 # (HETATM records) so add it here
14 _AT_HET_CA = IMP.atom.AtomType("HET: CA ")
15 
16 
17 def resnums2str(res):
18  """Take iterable of TempResidues and return compatified string"""
19  if len(res) == 0:
20  return ''
21  idxs = [r.get_index() for r in res]
22  idxs.sort()
23  all_ranges = []
24  cur_range = [idxs[0], idxs[0]]
25  for idx in idxs[1:]:
26  if idx != cur_range[1]+1:
27  all_ranges.append(cur_range)
28  cur_range = [idx, idx]
29  cur_range[1] = idx
30  all_ranges.append(cur_range)
31  ret = ''
32  for nr, r in enumerate(all_ranges):
33  ret += '%i-%i' % (r[0], r[1])
34  if nr < len(all_ranges)-1:
35  ret += ', '
36  return ret
37 
38 
39 def _select_ca_or_p(hiers, **kwargs):
40  """Select all CA (amino acids) or P (nucleic acids) as appropriate"""
41  sel_p = IMP.atom.Selection(hiers, atom_type=IMP.atom.AT_P, **kwargs)
42  ps = sel_p.get_selected_particles()
43  if ps:
44  # detected nucleotides. Selecting phosphorous instead of CA
45  return ps
46  else:
47  # Also select CA atoms in modified residues (such as MSE)
48  sel = IMP.atom.Selection(hiers, atom_type=IMP.atom.AT_CA, **kwargs) \
49  | IMP.atom.Selection(hiers, residue_type=IMP.atom.MSE,
50  atom_type=_AT_HET_CA, **kwargs)
51  return sel.get_selected_particles()
52 
53 
54 def get_structure(model, pdb_fn, chain_id, res_range=None, offset=0,
55  model_num=None, ca_only=False):
56  """read a structure from a PDB file and return a list of residues
57  @param model The IMP model
58  @param pdb_fn The file to read (in traditional PDB or mmCIF format)
59  @param chain_id Chain ID to read
60  @param res_range Add only a specific set of residues.
61  res_range[0] is the starting and res_range[1] is the ending
62  residue index
63  The ending residue can be "END", that will take everything
64  to the end of the sequence.
65  None gets you all.
66  @param offset Apply an offset to the residue indexes of the PDB file
67  @param model_num Read multi-model PDB and return that model (0-based index)
68  @param ca_only Read only CA atoms (by default, all non-waters are read)
69  """
70  # Read file in mmCIF format if requested
71  if pdb_fn.endswith('.cif'):
72  read_file = IMP.atom.read_mmcif
73  read_multi_file = IMP.atom.read_multimodel_mmcif
74  else:
75  read_file = IMP.atom.read_pdb
76  read_multi_file = IMP.atom.read_multimodel_pdb
77  if ca_only:
79  else:
81 
82  reader = read_file if model_num is None else read_multi_file
83  mh = reader(pdb_fn, model, IMP.atom.ChainPDBSelector([chain_id]) & sel)
84  if model_num is not None:
85  mh = mh[model_num]
86 
87  if res_range == [] or res_range is None:
88  ps = _select_ca_or_p(mh, chain=chain_id)
89  else:
90  start = res_range[0]
91  end = res_range[1]
92  if end == "END":
93  end = IMP.atom.Residue(
94  mh.get_children()[0].get_children()[-1]).get_index()
95  ps = _select_ca_or_p(mh, chain=chain_id,
96  residue_indexes=range(start, end+1))
97  ret = []
98 
99  for p in ps:
100  res = IMP.atom.Residue(IMP.atom.Atom(p).get_parent())
101  res.set_index(res.get_index() + offset)
102  ret.append(res)
103  if len(ret) == 0:
104  warnings.warn(
105  "no residues selected from %s in range %s" % (pdb_fn, res_range),
107  return ret
108 
109 
110 def build_bead(model, residues, input_coord=None):
111  """Generates a single bead"""
112 
113  ds_frag = (residues[0].get_index(), residues[-1].get_index())
114  prt = IMP.Particle(model)
116  ptem = IMP.core.XYZR(prt)
117  mass = IMP.atom.get_mass_from_number_of_residues(len(residues))
118 
119  if ds_frag[0] == ds_frag[-1]:
120  rt = residues[0].get_residue_type()
121  h = IMP.atom.Residue.setup_particle(prt, rt, ds_frag[0])
122  h.set_name('%i_bead' % (ds_frag[0]))
123  prt.set_name('%i_bead' % (ds_frag[0]))
124  try:
126  except IMP.ValueException:
128  IMP.atom.ResidueType("ALA"))
130  ptem.set_radius(radius)
131  else:
133  h.set_name('%i-%i_bead' % (ds_frag[0], ds_frag[-1]))
134  prt.set_name('%i-%i_bead' % (ds_frag[0], ds_frag[-1]))
135  h.set_residue_indexes(range(ds_frag[0], ds_frag[-1] + 1))
136  volume = IMP.atom.get_volume_from_mass(mass)
137  radius = 0.8 * (3.0 / 4.0 / pi * volume) ** (1.0 / 3.0)
138  ptem.set_radius(radius)
139 
141  try:
142  if tuple(input_coord) is not None:
143  ptem.set_coordinates(input_coord)
144  except TypeError:
145  pass
146  return h
147 
148 
149 def build_necklace(model, residues, resolution, input_coord=None):
150  """Generates a string of beads with given length"""
151  out_hiers = []
152  for chunk in list(IMP.pmi.tools.list_chunks_iterator(residues,
153  resolution)):
154  out_hiers.append(build_bead(model, chunk, input_coord=input_coord))
155  return out_hiers
156 
157 
158 def build_ca_centers(model, residues):
159  """Create a bead on the CA position with coarsened size and mass"""
160  out_hiers = []
161  for tempres in residues:
162  residue = tempres.get_hierarchy()
163  rp1 = IMP.Particle(model)
164  rp1.set_name("Residue_%i" % residue.get_index())
165  rt = residue.get_residue_type()
166  this_res = IMP.atom.Residue.setup_particle(rp1, residue)
167  try:
169  except IMP.ValueException:
171  IMP.atom.ResidueType("ALA"))
172  try:
173  mass = IMP.atom.get_mass(rt)
174  except Exception:
176  calpha = IMP.atom.Selection(
177  residue,
178  atom_types=[IMP.atom.AT_CA, _AT_HET_CA]).get_selected_particles()
179  cp = IMP.atom.Selection(
180  residue, atom_type=IMP.atom.AT_P).get_selected_particles()
181 
182  if len(calpha) == 1:
183  central_atom = calpha[0]
184  elif len(cp) == 1:
185  central_atom = cp[0]
186  else:
187  raise ValueError(
188  "build_ca_centers: weird selection (no CA, no "
189  "nucleotide P or ambiguous selection found)")
191  shape = IMP.algebra.Sphere3D(
192  IMP.core.XYZ(central_atom).get_coordinates(), radius)
193  IMP.core.XYZR.setup_particle(rp1, shape)
195  out_hiers.append(this_res)
196  return out_hiers
197 
198 
199 def setup_bead_as_gaussian(mh):
200  """Setup bead as spherical gaussian, using radius as variance"""
201  p = mh.get_particle()
202  center = IMP.core.XYZ(p).get_coordinates()
203  rad = IMP.core.XYZR(p).get_radius()
207  [rad]*3)
209 
210 
211 def show_representation(node):
212  print(node)
214  repr = IMP.atom.Representation(node)
215  resolutions = repr.get_resolutions()
216  for r in resolutions:
217  print('---- resolution %i ----' % r)
218  IMP.atom.show_molecular_hierarchy(repr.get_representation(r))
219  return True
220  else:
221  return False
222 
223 
224 def _get_color_for_representation(rep):
225  """Return an IMP.display.Color object (or None) for the given
226  Representation."""
227  if rep.color is not None:
228  if isinstance(rep.color, float):
229  return IMP.display.get_rgb_color(rep.color)
230  elif isinstance(rep.color, str):
231  return IMP.display.Color(*IMP.pmi.tools.color2rgb(rep.color))
232  elif hasattr(rep.color, '__iter__') and len(rep.color) == 3:
233  return IMP.display.Color(*rep.color)
234  elif isinstance(rep.color, IMP.display.Color):
235  return rep.color
236  else:
237  raise TypeError("Color must be Chimera color name, a hex "
238  "string, a float or (r,g,b) tuple")
239 
240 
241 def _add_fragment_provenance(fragment, first_residue, rephandler):
242  """Track the original source of a fragment's structure.
243  If the residues in the given fragment were extracted from a PDB
244  file, add suitable provenance information to the Model (the name
245  of that file, chain ID, and residue index offset)."""
246  pdb_element = rephandler.pdb_for_residue.get(first_residue.get_index())
247  if pdb_element:
248  m = fragment.get_model()
249  p = IMP.Particle(m, "input structure")
251  p, pdb_element.filename, pdb_element.chain_id, pdb_element.offset)
252  IMP.core.add_provenance(m, fragment, sp)
253  return pdb_element
254 
255 
256 def build_representation(parent, rep, coord_finder, rephandler):
257  """Create requested representation.
258  For beads, identifies continuous segments and sets up as Representation.
259  If any volume-based representations (e.g.,densities) are requested,
260  will instead create a single Representation node.
261  All reps are added as children of the passed parent.
262  @param parent The Molecule to which we'll add representations
263  @param rep What to build. An instance of pmi::topology::_Representation
264  @param coord_finder A _FindCloseStructure object to help localize beads
265  """
266  built_reps = []
267  atomic_res = 0
268  ca_res = 1
269  model = parent.hier.get_model()
270  color = _get_color_for_representation(rep)
271 
272  # first get the primary representation (currently, the smallest bead size)
273  # eventually we won't require beads to be present at all
274  primary_resolution = min(rep.bead_resolutions)
275 
276  # if collective densities, will return single node with everything
277  # below we sample or read the GMMs and add them as representation
278  # flag indicating grouping nonlinear segments with one GMM
279  single_node = False
280  prov_dict = {}
281  if rep.density_residues_per_component:
282  single_node = True
283  num_components = (len(rep.residues)
284  // rep.density_residues_per_component+1)
285  rep_dict = defaultdict(list)
286  segp = IMP.Particle(model)
287  root_representation = IMP.atom.Representation.setup_particle(
288  segp, primary_resolution)
289  built_reps.append(root_representation)
290  res_nums = [r.get_index() for r in rep.residues]
291  IMP.atom.Fragment.setup_particle(segp, res_nums)
292  density_frag = IMP.atom.Fragment.setup_particle(
293  IMP.Particle(model), res_nums)
294  density_frag.get_particle().set_name(
295  "Densities %i" % rep.density_residues_per_component)
296  density_ps = []
297 
298  if os.path.exists(rep.density_prefix + '.txt') \
299  and not rep.density_force_compute:
301  rep.density_prefix + '.txt', density_ps, model)
302  if (len(density_ps) != num_components
303  or not os.path.exists(rep.density_prefix + '.txt')
304  or rep.density_force_compute):
305  fit_coords = []
306  total_mass = 0.0
307  for r in rep.residues:
308  for p in IMP.core.get_leaves(r.hier):
309  fit_coords.append(IMP.core.XYZ(p).get_coordinates())
310  total_mass += IMP.atom.Mass(p).get_mass()
311 
312  # fit GMM
313  density_ps = []
315  num_components,
316  model,
317  density_ps,
318  min_covar=4.0,
319  mass_multiplier=total_mass)
320 
322  rep.density_prefix + '.txt')
323  if rep.density_voxel_size > 0.0:
325  density_ps, rep.density_prefix + '.mrc',
326  rep.density_voxel_size, fast=True)
327 
328  for n, d in enumerate(density_ps):
329  d.set_name('Density #%d' % n)
330  density_frag.add_child(d)
331  root_representation.add_representation(
332  density_frag, IMP.atom.DENSITIES,
333  rep.density_residues_per_component)
334 
335  # get continuous segments from residues
336  segments = []
337  rsort = sorted(list(rep.residues), key=lambda r: r.get_index())
338  prev_idx = rsort[0].get_index()-1
339  prev_structure = rsort[0].get_has_structure()
340  cur_seg = []
341  force_break = False
342  for nr, r in enumerate(rsort):
343  if (r.get_index() != prev_idx+1
344  or r.get_has_structure() != prev_structure or force_break):
345  segments.append(cur_seg)
346  cur_seg = []
347  force_break = False
348  cur_seg.append(r)
349  prev_idx = r.get_index()
350  prev_structure = r.get_has_structure()
351  if r.get_index()-1 in rep.bead_extra_breaks:
352  force_break = True
353  if cur_seg != []:
354  segments.append(cur_seg)
355 
356  # for each segment, merge into beads
357  name_all = 'frags:'
358  name_count = 0
359  for frag_res in segments:
360  res_nums = [r.get_index() for r in frag_res]
361  rrange = "%i-%i" % (res_nums[0], res_nums[-1])
362  name = "Frag_" + rrange
363  if name_count < 3:
364  name_all += rrange + ','
365  elif name_count == 3:
366  name_all += '...'
367  name_count += 1
368  segp = IMP.Particle(model, name)
369  IMP.atom.Fragment.setup_particle(segp, res_nums)
370  if not single_node:
371  this_representation = IMP.atom.Representation.setup_particle(
372  segp, primary_resolution)
373  built_reps.append(this_representation)
374  for resolution in rep.bead_resolutions:
375  fp = IMP.Particle(model)
376  this_resolution = IMP.atom.Fragment.setup_particle(fp, res_nums)
377  this_resolution.set_name("%s: Res %i" % (name, resolution))
378  if frag_res[0].get_has_structure():
379  pdb_element = _add_fragment_provenance(
380  this_resolution, frag_res[0], rephandler)
381  if pdb_element is not None:
382  prov_dict[resolution] = pdb_element
383  # if structured, merge particles as needed
384  if resolution == atomic_res:
385  for residue in frag_res:
386  this_resolution.add_child(residue.get_hierarchy())
387  elif resolution == ca_res and rep.bead_ca_centers:
388  beads = build_ca_centers(model, frag_res)
389  for bead in beads:
390  this_resolution.add_child(bead)
391  else:
393  "X")
394  for residue in frag_res:
395  tempc.add_child(IMP.atom.create_clone(residue.hier))
397  tempc, resolution)
398  for bead in beads.get_children():
399  this_resolution.add_child(bead)
400  del tempc
401  del beads
402  else:
403  # if unstructured, create necklace
404  input_coord = coord_finder.find_nearest_coord(
405  min(r.get_index() for r in frag_res))
406  if input_coord is None:
407  input_coord = rep.bead_default_coord
408  beads = build_necklace(model,
409  frag_res,
410  resolution,
411  input_coord)
412  for bead in beads:
413  this_resolution.add_child(bead)
414 
415  # if requested, color all resolutions the same
416  if color:
417  for lv in IMP.core.get_leaves(this_resolution):
419 
420  # finally decide where to put this resolution
421  # if volumetric, collect resolutions from different
422  # segments together
423  if single_node:
424  rep_dict[resolution] += this_resolution.get_children()
425  else:
426  if resolution == primary_resolution:
427  this_representation.add_child(this_resolution)
428  else:
429  this_representation.add_representation(this_resolution,
430  IMP.atom.BALLS,
431  resolution)
432  # if individual beads to be setup as Gaussians:
433  if rep.setup_particles_as_densities:
434  for p in IMP.core.get_leaves(this_resolution):
435  setup_bead_as_gaussian(p)
436  this_resolution.set_name(
437  this_resolution.get_name() + ' Densities %i' % resolution)
438  this_representation.add_representation(this_resolution,
439  IMP.atom.DENSITIES,
440  resolution)
441 
442  if single_node:
443  root_representation.set_name(name_all.strip(',') + ": Base")
444  d = root_representation.get_representations(IMP.atom.DENSITIES)
445  d[0].set_name('%s: ' % name_all + d[0].get_name())
446  for resolution in rep.bead_resolutions:
447  this_resolution = IMP.atom.Fragment.setup_particle(
448  IMP.Particle(model),
449  [r.get_index() for r in rep.residues])
450  this_resolution.set_name("%s: Res %i" % (name_all, resolution))
451  # Use provenance information from the last original node (hopefully
452  # all nodes have the same provenance, i.e. came from the same
453  # PDB file)
454  if prov_dict.get(resolution):
455  pdb_element = prov_dict[resolution]
457  IMP.Particle(model, "input structure"),
458  pdb_element.filename,
459  pdb_element.chain_id, pdb_element.offset)
460  IMP.core.add_provenance(model, this_resolution, sp)
461  for hier in rep_dict[resolution]:
462  this_resolution.add_child(hier)
463  if resolution == primary_resolution:
464  root_representation.add_child(this_resolution)
465  else:
466  root_representation.add_representation(this_resolution,
467  IMP.atom.BALLS,
468  resolution)
469  return built_reps
def list_chunks_iterator
Yield successive length-sized chunks from a list.
Definition: tools.py:603
Tools for handling Gaussian Mixture Models.
Definition: gmm_tools.py:1
Add mass to a particle.
Definition: Mass.h:23
double get_volume_from_residue_type(ResidueType rt)
Return an estimate for the volume of a given residue.
Simple 3D transformation class.
Represent an RGB color.
Definition: Color.h:25
static Gaussian setup_particle(Model *m, ParticleIndex pi)
Definition: core/Gaussian.h:65
void show_molecular_hierarchy(Hierarchy h)
Print out the molecular hierarchy.
static Fragment setup_particle(Model *m, ParticleIndex pi)
Definition: Fragment.h:67
double get_mass(const Selection &s)
Get the total mass of a hierarchy, in Daltons.
static XYZR setup_particle(Model *m, ParticleIndex pi)
Definition: XYZR.h:48
double get_mass(ResidueType c)
Get the mass from the residue type.
static StructureProvenance setup_particle(Model *m, ParticleIndex pi, std::string filename, std::string chain_id, int residue_offset)
Definition: provenance.h:157
Color get_rgb_color(double f)
Return the color for f from the RGB color map.
double get_mass_from_number_of_residues(unsigned int num_aa)
Estimate the mass of a protein from the number of amino acids.
Miscellaneous utilities.
Definition: tools.py:1
double get_ball_radius_from_volume_3d(double volume)
Return the radius of a sphere with a given volume.
Definition: Sphere3D.h:35
The type of an atom.
static Residue setup_particle(Model *m, ParticleIndex pi, ResidueType t, int index, int insertion_code)
Definition: Residue.h:160
static Representation setup_particle(Model *m, ParticleIndex pi)
GenericHierarchies get_leaves(Hierarchy mhd)
Get all the leaves of the bit of hierarchy.
A reference frame in 3D.
def color2rgb
Given a Chimera color name or hex color value, return RGB.
Definition: tools.py:1574
Warning related to handling of structures.
A Gaussian distribution in 3D.
Definition: Gaussian3D.h:25
def fit_gmm_to_points
fit a GMM to some points.
Definition: gmm_tools.py:243
A decorator for a representation.
double get_volume_from_mass(double m, ProteinDensityReference ref=ALBER)
Estimate the volume of a protein from its mass.
Ints get_index(const ParticlesTemp &particles, const Subset &subset, const Subsets &excluded)
A decorator for a particle representing an atom.
Definition: atom/Atom.h:238
static Mass setup_particle(Model *m, ParticleIndex pi, Float mass)
Definition: Mass.h:48
The type for a residue.
PDBSelector * get_default_pdb_selector()
Definition: pdb.h:559
A decorator for a particle with x,y,z coordinates.
Definition: XYZ.h:30
static Colored setup_particle(Model *m, ParticleIndex pi, Color color)
Definition: Colored.h:62
def write_gmm_to_map
write density map from GMM.
Definition: gmm_tools.py:118
A decorator for a residue.
Definition: Residue.h:137
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Hierarchy create_simplified_along_backbone(Chain input, const IntRanges &residue_segments, bool keep_detailed=false)
Rotation3D get_identity_rotation_3d()
Return a rotation that does not do anything.
Definition: Rotation3D.h:351
Class to handle individual particles of a Model object.
Definition: Particle.h:43
Select all CA ATOM records.
Definition: pdb.h:142
Python classes to represent, score, sample and analyze models.
def write_gmm_to_text
write a list of gaussians to text.
Definition: gmm_tools.py:61
Functionality for loading, creating, manipulating and scoring atomic structures.
void add_provenance(Model *m, ParticleIndex pi, Provenance p)
Add provenance to part of the model.
static Chain setup_particle(Model *m, ParticleIndex pi, std::string id)
Definition: Chain.h:83
An exception for an invalid value being passed to IMP.
Definition: exception.h:136
Select hierarchy particles identified by the biological name.
Definition: Selection.h:70
Select all ATOM and HETATM records with the given chain ids.
Definition: pdb.h:256
def decorate_gmm_from_text
read the output from write_gmm_to_text, decorate as Gaussian and Mass
Definition: gmm_tools.py:23
A decorator for a particle with x,y,z coordinates and a radius.
Definition: XYZR.h:27