IMP logo
IMP Reference Guide  2.15.0
The Integrative Modeling Platform
tools.py
1 #!/usr/bin/env python
2 
3 """@namespace IMP.pmi.tools
4  Miscellaneous utilities.
5 """
6 
7 from __future__ import print_function, division
8 import IMP
9 import IMP.algebra
10 import IMP.isd
11 import IMP.pmi
12 import IMP.pmi.topology
13 import collections
14 import itertools
15 import math
16 import sys
17 import ast
18 try:
19  from time import process_time # needs python 3.3 or later
20 except ImportError:
21  from time import clock as process_time
22 import RMF
23 import IMP.rmf
24 from collections import defaultdict, OrderedDict
25 import warnings
26 
27 
28 def _get_system_for_hier(hier):
29  """Given a hierarchy, return the System that created it, or None"""
30  # If we are given the raw particle, get the corresponding Hierarchy
31  # decorator if available
32  if hier and not hasattr(hier, 'get_parent'):
34  hier = IMP.atom.Hierarchy(hier)
35  else:
36  return None
37  while hier:
38  # See if we labeled the Python object directly with the System
39  if hasattr(hier, '_pmi2_system'):
40  return hier._pmi2_system()
41  # Otherwise (maybe we got a new Python wrapper around the same C++
42  # object), try all extant systems
43  for ws in IMP.pmi.topology.System._all_systems:
44  s = ws()
45  if s and s.hier == hier:
46  return s
47  # Try the next level up in the hierarchy
48  hier = hier.get_parent()
49 
50 
51 def _all_protocol_outputs(hier):
52  """Iterate over all (ProtocolOutput, State) pairs for the
53  given hierarchy"""
54  system = _get_system_for_hier(hier)
55  if system:
56  for state in system.states:
57  for p in state._protocol_output:
58  yield p
59 
60 
61 def _add_pmi_provenance(p):
62  """Tag the given particle as being created by the current version
63  of PMI."""
66  p, name="IMP PMI module", version=IMP.pmi.get_module_version(),
67  location="https://integrativemodeling.org")
69 
70 
71 def _get_restraint_set_keys():
72  if not hasattr(_get_restraint_set_keys, 'pmi_rs_key'):
73  _get_restraint_set_keys.pmi_rs_key = IMP.ModelKey("PMI restraints")
74  _get_restraint_set_keys.rmf_rs_key = IMP.ModelKey("RMF restraints")
75  return (_get_restraint_set_keys.pmi_rs_key,
76  _get_restraint_set_keys.rmf_rs_key)
77 
78 
79 def _add_restraint_sets(model, mk, mk_rmf):
80  rs = IMP.RestraintSet(model, "All PMI restraints")
81  rs_rmf = IMP.RestraintSet(model, "All PMI RMF restraints")
82  model.add_data(mk, rs)
83  model.add_data(mk_rmf, rs_rmf)
84  return rs, rs_rmf
85 
86 
87 def add_restraint_to_model(model, restraint, add_to_rmf=False):
88  """Add a PMI restraint to the model.
89  Since Model.add_restraint() no longer exists (in modern IMP restraints
90  should be added to a ScoringFunction instead) store them instead in
91  a RestraintSet, and keep a reference to it in the Model.
92 
93  If `add_to_rmf` is True, also add the restraint to a separate list
94  of restraints that will be written out to RMF files (by default, most
95  PMI restraints are not)."""
96  mk, mk_rmf = _get_restraint_set_keys()
97  if model.get_has_data(mk):
98  rs = IMP.RestraintSet.get_from(model.get_data(mk))
99  rs_rmf = IMP.RestraintSet.get_from(model.get_data(mk_rmf))
100  else:
101  rs, rs_rmf = _add_restraint_sets(model, mk, mk_rmf)
102  rs.add_restraint(restraint)
103  if add_to_rmf:
104  rs_rmf.add_restraint(restraint)
105 
106 
107 def get_restraint_set(model, rmf=False):
108  """Get a RestraintSet containing all PMI restraints added to the model.
109  If `rmf` is True, return only the subset of these restraints that
110  should be written out to RMF files."""
111  mk, mk_rmf = _get_restraint_set_keys()
112  if not model.get_has_data(mk):
113  warnings.warn("no restraints added to model yet",
115  _add_restraint_sets(model, mk, mk_rmf)
116  if rmf:
117  return IMP.RestraintSet.get_from(model.get_data(mk_rmf))
118  else:
119  return IMP.RestraintSet.get_from(model.get_data(mk))
120 
121 
122 class Stopwatch(object):
123  """Collect timing information.
124  Add an instance of this class to outputobjects to get timing information
125  in a stat file."""
126 
127  def __init__(self, isdelta=True):
128  """Constructor.
129  @param isdelta if True (the default) then report the time since the
130  last use of this class; if False, report cumulative time."""
131  self.starttime = process_time()
132  self.label = "None"
133  self.isdelta = isdelta
134 
135  def set_label(self, labelstr):
136  self.label = labelstr
137 
138  def get_output(self):
139  output = {}
140  if self.isdelta:
141  newtime = process_time()
142  output["Stopwatch_" + self.label + "_delta_seconds"] \
143  = str(newtime - self.starttime)
144  self.starttime = newtime
145  else:
146  output["Stopwatch_" + self.label + "_elapsed_seconds"] \
147  = str(process_time() - self.starttime)
148  return output
149 
150 
151 class SetupNuisance(object):
152 
153  def __init__(self, m, initialvalue, minvalue, maxvalue, isoptimized=True,
154  name=None):
155 
156  p = IMP.Particle(m)
157  if name:
158  p.set_name(name)
159  nuisance = IMP.isd.Scale.setup_particle(p, initialvalue)
160  if minvalue:
161  nuisance.set_lower(minvalue)
162  if maxvalue:
163  nuisance.set_upper(maxvalue)
164 
165  # m.add_score_state(IMP.core.SingletonConstraint(IMP.isd.NuisanceRangeModifier(),None,nuisance))
166  nuisance.set_is_optimized(nuisance.get_nuisance_key(), isoptimized)
167  self.nuisance = nuisance
168 
169  def get_particle(self):
170  return self.nuisance
171 
172 
173 class SetupWeight(object):
174 
175  def __init__(self, m, isoptimized=True, nweights_or_weights=None):
176  pw = IMP.Particle(m)
177  if isinstance(nweights_or_weights, int):
178  self.weight = IMP.isd.Weight.setup_particle(
179  pw, nweights_or_weights
180  )
181  else:
182  try:
183  nweights_or_weights = list(nweights_or_weights)
184  self.weight = IMP.isd.Weight.setup_particle(
185  pw, nweights_or_weights
186  )
187  except (TypeError, IMP.UsageException):
188  self.weight = IMP.isd.Weight.setup_particle(pw)
189  self.weight.set_weights_are_optimized(isoptimized)
190 
191  def get_particle(self):
192  return self.weight
193 
194 
195 class SetupSurface(object):
196 
197  def __init__(self, m, center, normal, isoptimized=True):
198  p = IMP.Particle(m)
199  self.surface = IMP.core.Surface.setup_particle(p, center, normal)
200  self.surface.set_coordinates_are_optimized(isoptimized)
201  self.surface.set_normal_is_optimized(isoptimized)
202 
203  def get_particle(self):
204  return self.surface
205 
206 
207 class ParticleToSampleList(object):
208 
209  def __init__(self, label="None"):
210 
211  self.dictionary_particle_type = {}
212  self.dictionary_particle_transformation = {}
213  self.dictionary_particle_name = {}
214  self.label = label
215 
216  def add_particle(self, particle, particle_type, particle_transformation,
217  name):
218  if particle_type not in ["Rigid_Bodies", "Floppy_Bodies", "Nuisances",
219  "X_coord", "Weights", "Surfaces"]:
220  raise TypeError("not the right particle type")
221  else:
222  self.dictionary_particle_type[particle] = particle_type
223  if particle_type == "Rigid_Bodies":
224  if (isinstance(particle_transformation, tuple)
225  and len(particle_transformation) == 2
226  and all(isinstance(x, float)
227  for x in particle_transformation)):
228  self.dictionary_particle_transformation[
229  particle] = particle_transformation
230  self.dictionary_particle_name[particle] = name
231  else:
232  raise TypeError(
233  "ParticleToSampleList: not the right transformation "
234  "format for Rigid_Bodies, should be a tuple of floats")
235  elif particle_type == "Surfaces":
236  if (isinstance(particle_transformation, tuple)
237  and len(particle_transformation) == 3
238  and all(isinstance(x, float)
239  for x in particle_transformation)):
240  self.dictionary_particle_transformation[
241  particle] = particle_transformation
242  self.dictionary_particle_name[particle] = name
243  else:
244  raise TypeError(
245  "ParticleToSampleList: not the right transformation "
246  "format for Surfaces, should be a tuple of floats")
247  else:
248  if isinstance(particle_transformation, float):
249  self.dictionary_particle_transformation[
250  particle] = particle_transformation
251  self.dictionary_particle_name[particle] = name
252  else:
253  raise TypeError(
254  "ParticleToSampleList: not the right transformation "
255  "format, should be a float")
256 
257  def get_particles_to_sample(self):
258  ps = {}
259  for particle in self.dictionary_particle_type:
260  key = self.dictionary_particle_type[particle] + \
261  "ParticleToSampleList_" + \
262  self.dictionary_particle_name[particle] + "_" + self.label
263  value = (
264  [particle],
265  self.dictionary_particle_transformation[particle])
266  ps[key] = value
267  return ps
268 
269 
270 def get_cross_link_data(directory, filename, dist, omega, sigma,
271  don=None, doff=None, prior=0, type_of_profile="gofr"):
272 
273  (distmin, distmax, ndist) = dist
274  (omegamin, omegamax, nomega) = omega
275  (sigmamin, sigmamax, nsigma) = sigma
276 
277  filen = IMP.isd.get_data_path("CrossLinkPMFs.dict")
278  with open(filen) as xlpot:
279  dictionary = ast.literal_eval(xlpot.readline())
280 
281  xpot = dictionary[directory][filename]["distance"]
282  pot = dictionary[directory][filename][type_of_profile]
283 
284  dist_grid = get_grid(distmin, distmax, ndist, False)
285  omega_grid = get_log_grid(omegamin, omegamax, nomega)
286  sigma_grid = get_log_grid(sigmamin, sigmamax, nsigma)
287 
288  if don is not None and doff is not None:
289  xlmsdata = IMP.isd.CrossLinkData(
290  dist_grid,
291  omega_grid,
292  sigma_grid,
293  xpot,
294  pot,
295  don,
296  doff,
297  prior)
298  else:
299  xlmsdata = IMP.isd.CrossLinkData(
300  dist_grid,
301  omega_grid,
302  sigma_grid,
303  xpot,
304  pot)
305  return xlmsdata
306 
307 
308 def get_grid(gmin, gmax, ngrid, boundaries):
309  grid = []
310  dx = (gmax - gmin) / float(ngrid)
311  for i in range(0, ngrid + 1):
312  if(not boundaries and i == 0):
313  continue
314  if(not boundaries and i == ngrid):
315  continue
316  grid.append(gmin + float(i) * dx)
317  return grid
318 
319 
320 def get_log_grid(gmin, gmax, ngrid):
321  grid = []
322  for i in range(0, ngrid + 1):
323  grid.append(gmin * math.exp(float(i) / ngrid * math.log(gmax / gmin)))
324  return grid
325 
326 
328  '''
329  example '"{ID_Score}" > 28 AND "{Sample}" ==
330  "%10_1%" OR ":Sample}" == "%10_2%" OR ":Sample}"
331  == "%10_3%" OR ":Sample}" == "%8_1%" OR ":Sample}" == "%8_2%"'
332  '''
333 
334  import pyparsing as pp
335 
336  operator = pp.Regex(">=|<=|!=|>|<|==|in").setName("operator")
337  value = pp.QuotedString(
338  '"') | pp.Regex(
339  r"[+-]?\d+(:?\.\d*)?(:?[eE][+-]?\d+)?")
340  identifier = pp.Word(pp.alphas, pp.alphanums + "_")
341  comparison_term = identifier | value
342  condition = pp.Group(comparison_term + operator + comparison_term)
343 
344  expr = pp.operatorPrecedence(condition, [
345  ("OR", 2, pp.opAssoc.LEFT, ),
346  ("AND", 2, pp.opAssoc.LEFT, ),
347  ])
348 
349  parsedstring = str(expr.parseString(inputstring)) \
350  .replace("[", "(") \
351  .replace("]", ")") \
352  .replace(",", " ") \
353  .replace("'", " ") \
354  .replace("%", "'") \
355  .replace("{", "float(entry['") \
356  .replace("}", "'])") \
357  .replace(":", "str(entry['") \
358  .replace("}", "'])") \
359  .replace("AND", "and") \
360  .replace("OR", "or")
361  return parsedstring
362 
363 
364 def open_file_or_inline_text(filename):
365  try:
366  fl = open(filename, "r")
367  except IOError:
368  fl = filename.split("\n")
369  return fl
370 
371 
372 def get_ids_from_fasta_file(fastafile):
373  ids = []
374  with open(fastafile) as ff:
375  for line in ff:
376  if line[0] == ">":
377  ids.append(line[1:-1])
378  return ids
379 
380 
381 def get_closest_residue_position(hier, resindex, terminus="N"):
382  '''
383  this function works with plain hierarchies, as read from the pdb,
384  no multi-scale hierarchies
385  '''
386  p = []
387  niter = 0
388  while len(p) == 0:
389  niter += 1
390  sel = IMP.atom.Selection(hier, residue_index=resindex,
391  atom_type=IMP.atom.AT_CA)
392 
393  if terminus == "N":
394  resindex += 1
395  if terminus == "C":
396  resindex -= 1
397 
398  if niter >= 10000:
399  print("get_closest_residue_position: exiting while loop "
400  "without result")
401  break
402  p = sel.get_selected_particles()
403 
404  if len(p) == 1:
405  return IMP.core.XYZ(p[0]).get_coordinates()
406  elif len(p) == 0:
407  print("get_closest_residue_position: got NO residues for hierarchy "
408  "%s and residue %i" % (hier, resindex))
409  raise Exception(
410  "get_closest_residue_position: got NO residues for hierarchy "
411  "%s and residue %i" % (hier, resindex))
412  else:
413  raise ValueError(
414  "got multiple residues for hierarchy %s and residue %i; the list "
415  "of particles is %s"
416  % (hier, resindex, str([pp.get_name() for pp in p])))
417 
418 
419 def get_residue_gaps_in_hierarchy(hierarchy, start, end):
420  '''
421  Return the residue index gaps and contiguous segments in the hierarchy.
422 
423  @param hierarchy hierarchy to examine
424  @param start first residue index
425  @param end last residue index
426 
427  @return A list of lists of the form
428  [[1,100,"cont"],[101,120,"gap"],[121,200,"cont"]]
429  '''
430  gaps = []
431  for n, rindex in enumerate(range(start, end + 1)):
432  sel = IMP.atom.Selection(hierarchy, residue_index=rindex,
433  atom_type=IMP.atom.AT_CA)
434 
435  if len(sel.get_selected_particles()) == 0:
436  if n == 0:
437  # set the initial condition
438  rindexgap = start
439  rindexcont = start - 1
440  if rindexgap == rindex - 1:
441  # residue is contiguous with the previously discovered gap
442  gaps[-1][1] += 1
443  else:
444  # residue is not contiguous with the previously discovered gap
445  # hence create a new gap tuple
446  gaps.append([rindex, rindex, "gap"])
447  # update the index of the last residue gap
448  rindexgap = rindex
449  else:
450  if n == 0:
451  # set the initial condition
452  rindexgap = start - 1
453  rindexcont = start
454  if rindexcont == rindex - 1:
455  # residue is contiguous with the previously discovered
456  # continuous part
457  gaps[-1][1] += 1
458  else:
459  # residue is not contiguous with the previously discovered
460  # continuous part, hence create a new cont tuple
461  gaps.append([rindex, rindex, "cont"])
462  # update the index of the last residue gap
463  rindexcont = rindex
464  return gaps
465 
466 
467 class map(object):
468 
469  def __init__(self):
470  self.map = {}
471 
472  def set_map_element(self, xvalue, yvalue):
473  self.map[xvalue] = yvalue
474 
475  def get_map_element(self, invalue):
476  if isinstance(invalue, float):
477  n = 0
478  mindist = 1
479  for x in self.map:
480  dist = (invalue - x) * (invalue - x)
481 
482  if n == 0:
483  mindist = dist
484  minx = x
485  if dist < mindist:
486  mindist = dist
487  minx = x
488  n += 1
489  return self.map[minx]
490  elif isinstance(invalue, str):
491  return self.map[invalue]
492  else:
493  raise TypeError("wrong type for map")
494 
495 
496 def select_by_tuple_2(hier, tuple_selection, resolution):
497  """New tuple format: molname OR (start,stop,molname,copynum,statenum)
498  Copy and state are optional. Can also use 'None' for them which will
499  get all. You can also pass -1 for stop which will go to the end.
500  Returns the particles
501  """
502  kwds = {} # going to accumulate keywords
503  kwds['resolution'] = resolution
504  if isinstance(tuple_selection, str):
505  kwds['molecule'] = tuple_selection
506  elif isinstance(tuple_selection, tuple):
507  rbegin = tuple_selection[0]
508  rend = tuple_selection[1]
509  kwds['molecule'] = tuple_selection[2]
510  try:
511  copynum = tuple_selection[3]
512  if copynum is not None:
513  kwds['copy_index'] = copynum
514  except: # noqa: E722
515  pass
516  try:
517  statenum = tuple_selection[4]
518  if statenum is not None:
519  kwds['state_index'] = statenum
520  except: # noqa: E722
521  pass
522  if rend == -1:
523  if rbegin > 1:
524  s = IMP.atom.Selection(hier, **kwds)
525  s -= IMP.atom.Selection(hier,
526  residue_indexes=range(1, rbegin),
527  **kwds)
528  return s.get_selected_particles()
529  else:
530  kwds['residue_indexes'] = range(rbegin, rend+1)
531  s = IMP.atom.Selection(hier, **kwds)
532  return s.get_selected_particles()
533 
534 
535 def get_db_from_csv(csvfilename):
536  import csv
537  outputlist = []
538  with open(csvfilename) as fh:
539  csvr = csv.DictReader(fh)
540  for ls in csvr:
541  outputlist.append(ls)
542  return outputlist
543 
544 
545 def get_prot_name_from_particle(p, list_of_names):
546  '''Return the component name provided a particle and a list of names'''
547  root = p
548  protname = root.get_name()
549  is_a_bead = False
550  while protname not in list_of_names:
551  root0 = root.get_parent()
552  if root0 == IMP.atom.Hierarchy():
553  return (None, None)
554  protname = root0.get_name()
555 
556  # check if that is a bead
557  # this piece of code might be dangerous if
558  # the hierarchy was called Bead :)
559  if "Beads" in protname:
560  is_a_bead = True
561  root = root0
562  return (protname, is_a_bead)
563 
564 
566  '''
567  Retrieve the residue indexes for the given particle.
568 
569  The particle must be an instance of Fragment,Residue, Atom or Molecule
570  or else returns an empty list
571  '''
572  resind = []
574  resind = IMP.atom.Fragment(hier).get_residue_indexes()
576  resind = [IMP.atom.Residue(hier).get_index()]
577  elif IMP.atom.Atom.get_is_setup(hier):
578  a = IMP.atom.Atom(hier)
579  resind = [IMP.atom.Residue(a.get_parent()).get_index()]
581  resind_tmp = IMP.pmi.tools.OrderedSet()
582  for lv in IMP.atom.get_leaves(hier):
586  for ind in get_residue_indexes(lv):
587  resind_tmp.add(ind)
588  resind = list(resind_tmp)
589  else:
590  resind = []
591  return resind
592 
593 
594 def sort_by_residues(particles):
595  particles_residues = [(p, IMP.pmi.tools.get_residue_indexes(p))
596  for p in particles]
597  sorted_particles_residues = sorted(
598  particles_residues,
599  key=lambda tup: tup[1])
600  particles = [p[0] for p in sorted_particles_residues]
601  return particles
602 
603 #
604 # Parallel Computation
605 #
606 
607 
609  """Synchronize data over a parallel run"""
610  from mpi4py import MPI
611  comm = MPI.COMM_WORLD
612  rank = comm.Get_rank()
613  number_of_processes = comm.size
614  comm.Barrier()
615  if rank != 0:
616  comm.send(data, dest=0, tag=11)
617 
618  elif rank == 0:
619  for i in range(1, number_of_processes):
620  data_tmp = comm.recv(source=i, tag=11)
621  if isinstance(data, list):
622  data += data_tmp
623  elif isinstance(data, dict):
624  data.update(data_tmp)
625  else:
626  raise TypeError("data not supported, use list or dictionaries")
627 
628  for i in range(1, number_of_processes):
629  comm.send(data, dest=i, tag=11)
630 
631  if rank != 0:
632  data = comm.recv(source=0, tag=11)
633  return data
634 
635 #
636 # Lists and iterators
637 #
638 
639 
640 def sublist_iterator(ls, lmin=None, lmax=None):
641  '''
642  Yield all sublists of length >= lmin and <= lmax
643  '''
644  if lmin is None:
645  lmin = 0
646  if lmax is None:
647  lmax = len(ls)
648  n = len(ls) + 1
649  for i in range(n):
650  for j in range(i + 1, n):
651  if len(ls[i:j]) <= lmax and len(ls[i:j]) >= lmin:
652  yield ls[i:j]
653 
654 
655 def flatten_list(ls):
656  return [item for sublist in ls for item in sublist]
657 
658 
659 def list_chunks_iterator(list, length):
660  """ Yield successive length-sized chunks from a list.
661  """
662  for i in range(0, len(list), length):
663  yield list[i:i + length]
664 
665 
666 def chunk_list_into_segments(seq, num):
667  seq = list(seq)
668  avg = len(seq) / float(num)
669  out = []
670  last = 0.0
671 
672  while last < len(seq):
673  out.append(seq[int(last):int(last + avg)])
674  last += avg
675 
676  return out
677 
678 
679 class Segments(object):
680 
681  ''' This class stores integers
682  in ordered compact lists eg:
683  [[1,2,3],[6,7,8]]
684  the methods help splitting and merging the internal lists
685  Example:
686  s=Segments([1,2,3]) is [[1,2,3]]
687  s.add(4) is [[1,2,3,4]] (add right)
688  s.add(3) is [[1,2,3,4]] (item already existing)
689  s.add(7) is [[1,2,3,4],[7]] (new list)
690  s.add([8,9]) is [[1,2,3,4],[7,8,9]] (add item right)
691  s.add([5,6]) is [[1,2,3,4,5,6,7,8,9]] (merge)
692  s.remove(3) is [[1,2],[4,5,6,7,8,9]] (split)
693  etc.
694  '''
695 
696  def __init__(self, index):
697  '''index can be a integer or a list of integers '''
698  if isinstance(index, int):
699  self.segs = [[index]]
700  elif isinstance(index, list):
701  self.segs = [[index[0]]]
702  for i in index[1:]:
703  self.add(i)
704  else:
705  raise TypeError("index must be an int or list of ints")
706 
707  def add(self, index):
708  '''index can be a integer or a list of integers '''
709  if isinstance(index, int):
710  mergeleft = None
711  mergeright = None
712  for n, s in enumerate(self.segs):
713  if index in s:
714  return 0
715  else:
716  if s[0]-index == 1:
717  mergeleft = n
718  if index-s[-1] == 1:
719  mergeright = n
720  if mergeright is None and mergeleft is None:
721  self.segs.append([index])
722  if mergeright is not None and mergeleft is None:
723  self.segs[mergeright].append(index)
724  if mergeleft is not None and mergeright is None:
725  self.segs[mergeleft] = [index]+self.segs[mergeleft]
726  if mergeleft is not None and mergeright is not None:
727  self.segs[mergeright] = \
728  self.segs[mergeright]+[index]+self.segs[mergeleft]
729  del self.segs[mergeleft]
730 
731  for n in range(len(self.segs)):
732  self.segs[n].sort()
733 
734  self.segs.sort(key=lambda tup: tup[0])
735 
736  elif isinstance(index, list):
737  for i in index:
738  self.add(i)
739  else:
740  raise TypeError("index must be an int or list of ints")
741 
742  def remove(self, index):
743  '''index can be a integer'''
744  for n, s in enumerate(self.segs):
745  if index in s:
746  if s[0] == index:
747  self.segs[n] = s[1:]
748  elif s[-1] == index:
749  self.segs[n] = s[:-1]
750  else:
751  i = self.segs[n].index(index)
752  self.segs[n] = s[:i]
753  self.segs.append(s[i+1:])
754  for n in range(len(self.segs)):
755  self.segs[n].sort()
756  if len(self.segs[n]) == 0:
757  del self.segs[n]
758  self.segs.sort(key=lambda tup: tup[0])
759 
760  def get_flatten(self):
761  ''' Returns a flatten list '''
762  return [item for sublist in self.segs for item in sublist]
763 
764  def __repr__(self):
765  ret_tmp = "["
766  for seg in self.segs:
767  ret_tmp += str(seg[0])+"-"+str(seg[-1])+","
768  ret = ret_tmp[:-1]+"]"
769  return ret
770 
771 #
772 # Tools to simulate data
773 #
774 
775 
776 def normal_density_function(expected_value, sigma, x):
777  return (
778  1 / math.sqrt(2 * math.pi) / sigma *
779  math.exp(-(x - expected_value) ** 2 / 2 / sigma / sigma)
780  )
781 
782 
783 def log_normal_density_function(expected_value, sigma, x):
784  return (
785  1 / math.sqrt(2 * math.pi) / sigma / x *
786  math.exp(-(math.log(x / expected_value) ** 2 / 2 / sigma / sigma))
787  )
788 
789 
790 def print_multicolumn(list_of_strings, ncolumns=2, truncate=40):
791 
792  ls = list_of_strings
793 
794  cols = ncolumns
795  # add empty entries after ls
796  for i in range(len(ls) % cols):
797  ls.append(" ")
798 
799  split = [ls[i:i + len(ls) // cols]
800  for i in range(0, len(ls), len(ls) // cols)]
801  for row in zip(*split):
802  print("".join(str.ljust(i, truncate) for i in row))
803 
804 
805 class ColorChange(object):
806  '''Change color code to hexadecimal to rgb'''
807  def __init__(self):
808  self._NUMERALS = '0123456789abcdefABCDEF'
809  self._HEXDEC = dict((v, int(v, 16)) for v in
810  (x+y for x in self._NUMERALS
811  for y in self._NUMERALS))
812  self.LOWERCASE, self.UPPERCASE = 'x', 'X'
813 
814  def rgb(self, triplet):
815  return (float(self._HEXDEC[triplet[0:2]]),
816  float(self._HEXDEC[triplet[2:4]]),
817  float(self._HEXDEC[triplet[4:6]]))
818 
819  def triplet(self, rgb, lettercase=None):
820  if lettercase is None:
821  lettercase = self.LOWERCASE
822  return format(rgb[0] << 16 | rgb[1] << 8 | rgb[2], '06'+lettercase)
823 
824 
825 # -------------- Collections --------------- #
826 class OrderedSet(collections.MutableSet):
827 
828  def __init__(self, iterable=None):
829  self.end = end = []
830  end += [None, end, end] # sentinel node for doubly linked list
831  self.map = {} # key --> [key, prev, next]
832  if iterable is not None:
833  self |= iterable
834 
835  def __len__(self):
836  return len(self.map)
837 
838  def __contains__(self, key):
839  return key in self.map
840 
841  def add(self, key):
842  if key not in self.map:
843  end = self.end
844  curr = end[1]
845  curr[2] = end[1] = self.map[key] = [key, curr, end]
846 
847  def discard(self, key):
848  if key in self.map:
849  key, prev, next = self.map.pop(key)
850  prev[2] = next
851  next[1] = prev
852 
853  def __iter__(self):
854  end = self.end
855  curr = end[2]
856  while curr is not end:
857  yield curr[0]
858  curr = curr[2]
859 
860  def __reversed__(self):
861  end = self.end
862  curr = end[1]
863  while curr is not end:
864  yield curr[0]
865  curr = curr[1]
866 
867  def pop(self, last=True):
868  if not self:
869  raise KeyError('set is empty')
870  if last:
871  key = self.end[1][0]
872  else:
873  key = self.end[2][0]
874  self.discard(key)
875  return key
876 
877  def __repr__(self):
878  if not self:
879  return '%s()' % (self.__class__.__name__,)
880  return '%s(%r)' % (self.__class__.__name__, list(self))
881 
882  def __eq__(self, other):
883  if isinstance(other, OrderedSet):
884  return len(self) == len(other) and list(self) == list(other)
885  return set(self) == set(other)
886 
887 
888 class OrderedDefaultDict(OrderedDict):
889  """Store objects in order they were added, but with default type.
890  Source: http://stackoverflow.com/a/4127426/2608793
891  """
892  def __init__(self, *args, **kwargs):
893  if not args:
894  self.default_factory = None
895  else:
896  if not (args[0] is None or callable(args[0])):
897  raise TypeError('first argument must be callable or None')
898  self.default_factory = args[0]
899  args = args[1:]
900  super(OrderedDefaultDict, self).__init__(*args, **kwargs)
901 
902  def __missing__(self, key):
903  if self.default_factory is None:
904  raise KeyError(key)
905  self[key] = default = self.default_factory()
906  return default
907 
908  def __reduce__(self): # optional, for pickle support
909  args = (self.default_factory,) if self.default_factory else ()
910  if sys.version_info[0] >= 3:
911  return self.__class__, args, None, None, self.items()
912  else:
913  return self.__class__, args, None, None, self.iteritems()
914 
915 
916 # -------------- PMI2 Tools --------------- #
917 
918 def set_coordinates_from_rmf(hier, rmf_fn, frame_num=0):
919  """Extract frame from RMF file and fill coordinates. Must be identical
920  topology.
921 
922  @param hier The (System) hierarchy to fill (e.g. after you've built it)
923  @param rmf_fn The file to extract from
924  @param frame_num The frame number to extract
925  """
926  rh = RMF.open_rmf_file_read_only(rmf_fn)
927  IMP.rmf.link_hierarchies(rh, [hier])
928  IMP.rmf.load_frame(rh, RMF.FrameID(frame_num))
929  del rh
930 
931 
932 def input_adaptor(stuff, pmi_resolution=0, flatten=False, selection_tuple=None,
933  warn_about_slices=True):
934  """Adapt things for PMI (degrees of freedom, restraints, ...)
935  Returns list of list of hierarchies, separated into Molecules if possible.
936  The input can be a list, or a list of lists (iterable of ^1 or
937  iterable of ^2)
938  (iterable of ^2) Hierarchy -> returns input as list of list of hierarchies,
939  only one entry, not grouped by molecules.
940  (iterable of ^2) PMI::System/State/Molecule/TempResidue ->
941  returns residue hierarchies, grouped in molecules, at requested
942  resolution
943 
944  @param stuff Can be one of the following inputs:
945  IMP Hierarchy, PMI System/State/Molecule/TempResidue, or a
946  list/set (of list/set) of them.
947  Must be uniform input, however. No mixing object types.
948  @param pmi_resolution For selecting, only does it if you pass PMI
949  objects. Set it to "all" if you want all resolutions!
950  @param flatten Set to True if you just want all hierarchies in one list.
951  @param warn_about_slices Print a warning if you are requesting only part
952  of a bead. Sometimes you just don't care!
953  @note since this relies on IMP::atom::Selection, this will not return
954  any objects if they weren't built! But there should be no problem
955  if you request unbuilt residues - they should be ignored.
956  """
957 
958  if stuff is None:
959  return stuff
960 
961  if hasattr(stuff, '__iter__'):
962  if len(stuff) == 0:
963  return stuff
964  thelist = list(stuff)
965 
966  # iter of iter of should be ok
967  if all(hasattr(el, '__iter__') for el in thelist):
968  thelist = [i for sublist in thelist for i in sublist]
969  elif any(hasattr(el, '__iter__') for el in thelist):
970  raise Exception('input_adaptor: input_object must be a list '
971  'or a list of lists')
972 
973  stuff = thelist
974  else:
975  stuff = [stuff]
976 
977  # check that it is a hierarchy homogenously:
978  try:
979  is_hierarchy = all(IMP.atom.Hierarchy.get_is_setup(s) for s in stuff)
980  except (NotImplementedError, TypeError):
981  is_hierarchy = False
982  # get the other types homogenously
983  is_system = all(isinstance(s, IMP.pmi.topology.System) for s in stuff)
984  is_state = all(isinstance(s, IMP.pmi.topology.State) for s in stuff)
985  is_molecule = all(isinstance(s, IMP.pmi.topology.Molecule) for s in stuff)
986  is_temp_residue = all(isinstance(s, IMP.pmi.topology.TempResidue)
987  for s in stuff)
988 
989  # now that things are ok, do selection if requested
990  hier_list = []
991  pmi_input = False
992  if is_system or is_state or is_molecule or is_temp_residue:
993  # if PMI, perform selection using gathered indexes
994  pmi_input = True
995  # key is Molecule object, value are residues
996  indexes_per_mol = OrderedDefaultDict(list)
997  if is_system:
998  for system in stuff:
999  for state in system.get_states():
1000  mdict = state.get_molecules()
1001  for molname in mdict:
1002  for copy in mdict[molname]:
1003  indexes_per_mol[copy] += \
1004  [r.get_index() for r in copy.get_residues()]
1005  elif is_state:
1006  for state in stuff:
1007  mdict = state.get_molecules()
1008  for molname in mdict:
1009  for copy in mdict[molname]:
1010  indexes_per_mol[copy] += [r.get_index()
1011  for r in copy.get_residues()]
1012  elif is_molecule:
1013  for molecule in stuff:
1014  indexes_per_mol[molecule] += [r.get_index()
1015  for r in molecule.get_residues()]
1016  else: # is_temp_residue
1017  for tempres in stuff:
1018  indexes_per_mol[tempres.get_molecule()].append(
1019  tempres.get_index())
1020  for mol in indexes_per_mol:
1021  if pmi_resolution == 'all':
1022  # because you select from the molecule,
1023  # this will start the search from the base resolution
1025  mol.get_hierarchy(), residue_indexes=indexes_per_mol[mol])
1026  else:
1027  sel = IMP.atom.Selection(mol.get_hierarchy(),
1028  resolution=pmi_resolution,
1029  residue_indexes=indexes_per_mol[mol])
1030  ps = sel.get_selected_particles()
1031 
1032  # check that you don't have any incomplete fragments!
1033  if warn_about_slices:
1034  rset = set(indexes_per_mol[mol])
1035  for p in ps:
1037  fset = set(IMP.atom.Fragment(p).get_residue_indexes())
1038  if not fset <= rset:
1039  minset = min(fset)
1040  maxset = max(fset)
1041  found = fset & rset
1042  minf = min(found)
1043  maxf = max(found)
1044  resbreak = maxf if minf == minset else minset-1
1045  warnings.warn(
1046  'You are trying to select only part of the '
1047  'bead %s:%i-%i. The residues you requested '
1048  'are %i-%i. You can fix this by: '
1049  '1) requesting the whole bead/none of it; or'
1050  '2) break the bead up by passing '
1051  'bead_extra_breaks=[\'%i\'] in '
1052  'molecule.add_representation()'
1053  % (mol.get_name(), minset, maxset, minf, maxf,
1054  resbreak), IMP.pmi.ParameterWarning)
1055  hier_list.append([IMP.atom.Hierarchy(p) for p in ps])
1056  elif is_hierarchy:
1057  # check
1058  ps = []
1059  if pmi_resolution == 'all':
1060  for h in stuff:
1061  ps += select_at_all_resolutions(h)
1062  else:
1063  for h in stuff:
1064  ps += IMP.atom.Selection(
1065  h, resolution=pmi_resolution).get_selected_particles()
1066  hier_list = [IMP.atom.Hierarchy(p) for p in ps]
1067  if not flatten:
1068  hier_list = [hier_list]
1069  else:
1070  raise Exception('input_adaptor: you passed something of wrong type '
1071  'or a list with mixed types')
1072 
1073  if flatten and pmi_input:
1074  return [h for sublist in hier_list for h in sublist]
1075  else:
1076  return hier_list
1077 
1078 
1080  """Returns sequence-sorted segments array, each containing the first
1081  particle the last particle and the first residue index."""
1082 
1083  from operator import itemgetter
1084  hiers = IMP.pmi.tools.input_adaptor(mol)
1085  if len(hiers) > 1:
1086  raise ValueError("only pass stuff from one Molecule, please")
1087  hiers = hiers[0]
1088  segs = []
1089  for h in hiers:
1090  try:
1091  start = IMP.atom.Hierarchy(h).get_children()[0]
1092  except IndexError:
1093  start = IMP.atom.Hierarchy(h)
1094 
1095  try:
1096  end = IMP.atom.Hierarchy(h).get_children()[-1]
1097  except IndexError:
1098  end = IMP.atom.Hierarchy(h)
1099 
1100  startres = IMP.pmi.tools.get_residue_indexes(start)[0]
1101  segs.append((start, end, startres))
1102  return sorted(segs, key=itemgetter(2))
1103 
1104 
1105 def display_bonds(mol):
1106  """Decorate the sequence-consecutive particles from a PMI2 molecule
1107  with a bond, so that they appear connected in the rmf file"""
1108  SortedSegments = get_sorted_segments(mol)
1109  for x in range(len(SortedSegments) - 1):
1110 
1111  last = SortedSegments[x][1]
1112  first = SortedSegments[x + 1][0]
1113 
1114  p1 = last.get_particle()
1115  p2 = first.get_particle()
1116  if not IMP.atom.Bonded.get_is_setup(p1):
1118  if not IMP.atom.Bonded.get_is_setup(p2):
1120 
1123  IMP.atom.Bonded(p1),
1124  IMP.atom.Bonded(p2), 1)
1125 
1126 
1127 def get_all_leaves(list_of_hs):
1128  """ Just get the leaves from a list of hierarchies """
1129  lvs = list(itertools.chain.from_iterable(
1130  IMP.atom.get_leaves(item) for item in list_of_hs))
1131  return lvs
1132 
1133 
1134 def select_at_all_resolutions(hier=None, hiers=None, **kwargs):
1135  """Perform selection using the usual keywords but return ALL
1136  resolutions (BEADS and GAUSSIANS).
1137  Returns in flat list!
1138  """
1139 
1140  if hiers is None:
1141  hiers = []
1142  if hier is not None:
1143  hiers.append(hier)
1144  if len(hiers) == 0:
1145  warnings.warn("You passed nothing to select_at_all_resolutions()",
1147  return []
1148  ret = OrderedSet()
1149  for hsel in hiers:
1150  try:
1151  htest = IMP.atom.Hierarchy.get_is_setup(hsel)
1152  except: # noqa: E722
1153  raise Exception('select_at_all_resolutions: you have to pass '
1154  'an IMP Hierarchy')
1155  if not htest:
1156  raise Exception('select_at_all_resolutions: you have to pass '
1157  'an IMP Hierarchy')
1158  if 'resolution' in kwargs or 'representation_type' in kwargs:
1159  raise Exception("don't pass resolution or representation_type "
1160  "to this function")
1161  selB = IMP.atom.Selection(hsel, resolution=IMP.atom.ALL_RESOLUTIONS,
1162  representation_type=IMP.atom.BALLS,
1163  **kwargs)
1164  selD = IMP.atom.Selection(hsel, resolution=IMP.atom.ALL_RESOLUTIONS,
1165  representation_type=IMP.atom.DENSITIES,
1166  **kwargs)
1167  ret |= OrderedSet(selB.get_selected_particles())
1168  ret |= OrderedSet(selD.get_selected_particles())
1169  return list(ret)
1170 
1171 
1173  target_ps,
1174  sel_zone,
1175  entire_residues,
1176  exclude_backbone):
1177  """Utility to retrieve particles from a hierarchy within a
1178  zone around a set of ps.
1179  @param hier The hierarchy in which to look for neighbors
1180  @param target_ps The particles for zoning
1181  @param sel_zone The maximum distance
1182  @param entire_residues If True, will grab entire residues
1183  @param exclude_backbone If True, will only return sidechain particles
1184  """
1185 
1186  test_sel = IMP.atom.Selection(hier)
1187  backbone_types = ['C', 'N', 'CB', 'O']
1188  if exclude_backbone:
1189  test_sel -= IMP.atom.Selection(
1190  hier, atom_types=[IMP.atom.AtomType(n) for n in backbone_types])
1191  test_ps = test_sel.get_selected_particles()
1192  nn = IMP.algebra.NearestNeighbor3D([IMP.core.XYZ(p).get_coordinates()
1193  for p in test_ps])
1194  zone = set()
1195  for target in target_ps:
1196  zone |= set(nn.get_in_ball(IMP.core.XYZ(target).get_coordinates(),
1197  sel_zone))
1198  zone_ps = [test_ps[z] for z in zone]
1199  if entire_residues:
1200  final_ps = set()
1201  for z in zone_ps:
1202  final_ps |= set(IMP.atom.Hierarchy(z).get_parent().get_children())
1203  zone_ps = [h.get_particle() for h in final_ps]
1204  return zone_ps
1205 
1206 
1208  """Returns unique objects in original order"""
1209  rbs = set()
1210  beads = []
1211  rbs_ordered = []
1212  if not hasattr(hiers, '__iter__'):
1213  hiers = [hiers]
1214  for p in get_all_leaves(hiers):
1216  rb = IMP.core.RigidMember(p).get_rigid_body()
1217  if rb not in rbs:
1218  rbs.add(rb)
1219  rbs_ordered.append(rb)
1221  rb = IMP.core.NonRigidMember(p).get_rigid_body()
1222  if rb not in rbs:
1223  rbs.add(rb)
1224  rbs_ordered.append(rb)
1225  beads.append(p)
1226  else:
1227  beads.append(p)
1228  return rbs_ordered, beads
1229 
1230 
1231 def get_molecules(input_objects):
1232  "This function returns the parent molecule hierarchies of given objects"
1233  stuff = input_adaptor(input_objects, pmi_resolution='all', flatten=True)
1234  molecules = set()
1235  for h in stuff:
1236  is_root = False
1237  is_molecule = False
1238  while not (is_molecule or is_root):
1239  root = IMP.atom.get_root(h)
1240  if root == h:
1241  is_root = True
1242  is_molecule = IMP.atom.Molecule.get_is_setup(h)
1243  if is_molecule:
1244  molecules.add(IMP.atom.Molecule(h))
1245  h = h.get_parent()
1246  return list(molecules)
1247 
1248 
1249 def get_molecules_dictionary(input_objects):
1250  moldict = defaultdict(list)
1251  for mol in IMP.pmi.tools.get_molecules(input_objects):
1252  name = mol.get_name()
1253  moldict[name].append(mol)
1254 
1255  for mol in moldict:
1256  moldict[mol].sort(key=lambda x: IMP.atom.Copy(x).get_copy_index())
1257  return moldict
1258 
1259 
1260 def get_molecules_dictionary_by_copy(input_objects):
1261  moldict = defaultdict(dict)
1262  for mol in IMP.pmi.tools.get_molecules(input_objects):
1263  name = mol.get_name()
1264  c = IMP.atom.Copy(mol).get_copy_index()
1265  moldict[name][c] = mol
1266  return moldict
1267 
1268 
1269 def get_selections_dictionary(input_objects):
1270  moldict = IMP.pmi.tools.get_molecules_dictionary(input_objects)
1271  seldict = defaultdict(list)
1272  for name, mols in moldict.items():
1273  for m in mols:
1274  seldict[name].append(IMP.atom.Selection(m))
1275  return seldict
1276 
1277 
1278 def get_densities(input_objects):
1279  """Given a list of PMI objects, returns all density hierarchies within
1280  these objects. The output of this function can be inputted into
1281  things such as EM restraints. This function is intended to gather
1282  density particles appended to molecules (and not other hierarchies
1283  which might have been appended to the root node directly).
1284  """
1285  # Note that Densities can only be selected at the Root or Molecule
1286  # level and not at the Leaves level.
1287  # we'll first get all molecule hierarchies corresponding to the leaves.
1288  molecules = get_molecules(input_objects)
1289  densities = []
1290  for i in molecules:
1291  densities += IMP.atom.Selection(
1292  i, representation_type=IMP.atom.DENSITIES).get_selected_particles()
1293  return densities
1294 
1295 
1296 def shuffle_configuration(objects,
1297  max_translation=300., max_rotation=2.0 * math.pi,
1298  avoidcollision_rb=True, avoidcollision_fb=False,
1299  cutoff=10.0, niterations=100,
1300  bounding_box=None,
1301  excluded_rigid_bodies=[],
1302  hierarchies_excluded_from_collision=[],
1303  hierarchies_included_in_collision=[],
1304  verbose=False,
1305  return_debug=False):
1306  """Shuffle particles. Used to restart the optimization.
1307  The configuration of the system is initialized by placing each
1308  rigid body and each bead randomly in a box. If `bounding_box` is
1309  specified, the particles are placed inside this box; otherwise, each
1310  particle is displaced by up to max_translation angstroms, and randomly
1311  rotated. Effort is made to place particles far enough from each other to
1312  prevent any steric clashes.
1313  @param objects Can be one of the following inputs:
1314  IMP Hierarchy, PMI System/State/Molecule/TempResidue, or
1315  a list/set of them
1316  @param max_translation Max translation (rbs and flexible beads)
1317  @param max_rotation Max rotation (rbs only)
1318  @param avoidcollision_rb check if the particle/rigid body was
1319  placed close to another particle; uses the optional
1320  arguments cutoff and niterations
1321  @param avoidcollision_fb Advanced. Generally you want this False because
1322  it's hard to shuffle beads.
1323  @param cutoff Distance less than this is a collision
1324  @param niterations How many times to try avoiding collision
1325  @param bounding_box Only shuffle particles within this box.
1326  Defined by ((x1,y1,z1),(x2,y2,z2)).
1327  @param excluded_rigid_bodies Don't shuffle these rigid body objects
1328  @param hierarchies_excluded_from_collision Don't count collision
1329  with these bodies
1330  @param hierarchies_included_in_collision Hierarchies that are not
1331  shuffled, but should be included in collision calculation
1332  (for fixed regions)
1333  @param verbose Give more output
1334  @note Best to only call this function after you've set up degrees
1335  of freedom
1336  For debugging purposes, returns: <shuffled indexes>,
1337  <collision avoided indexes>
1338  """
1339 
1340  # checking input
1341  hierarchies = IMP.pmi.tools.input_adaptor(objects,
1342  pmi_resolution='all',
1343  flatten=True)
1344  rigid_bodies, flexible_beads = get_rbs_and_beads(hierarchies)
1345  if len(rigid_bodies) > 0:
1346  mdl = rigid_bodies[0].get_model()
1347  elif len(flexible_beads) > 0:
1348  mdl = flexible_beads[0].get_model()
1349  else:
1350  raise Exception("Could not find any particles in the hierarchy")
1351  if len(rigid_bodies) == 0:
1352  print("shuffle_configuration: rigid bodies were not initialized")
1353 
1354  # gather all particles
1356  gcpf.set_distance(cutoff)
1357 
1358  # Add particles from excluded hierarchies to excluded list
1359  collision_excluded_hierarchies = IMP.pmi.tools.input_adaptor(
1360  hierarchies_excluded_from_collision, pmi_resolution='all',
1361  flatten=True)
1362 
1363  collision_included_hierarchies = IMP.pmi.tools.input_adaptor(
1364  hierarchies_included_in_collision, pmi_resolution='all', flatten=True)
1365 
1366  collision_excluded_idxs = set(
1367  leaf.get_particle().get_index()
1368  for h in collision_excluded_hierarchies
1369  for leaf in IMP.core.get_leaves(h))
1370 
1371  collision_included_idxs = set(
1372  leaf.get_particle().get_index()
1373  for h in collision_included_hierarchies
1374  for leaf in IMP.core.get_leaves(h))
1375 
1376  # Excluded collision with Gaussians
1377  all_idxs = [] # expand to representations?
1378  for p in IMP.pmi.tools.get_all_leaves(hierarchies):
1380  all_idxs.append(p.get_particle_index())
1382  collision_excluded_idxs.add(p.get_particle_index())
1383 
1384  if bounding_box is not None:
1385  ((x1, y1, z1), (x2, y2, z2)) = bounding_box
1386  ub = IMP.algebra.Vector3D(x1, y1, z1)
1387  lb = IMP.algebra.Vector3D(x2, y2, z2)
1388  bb = IMP.algebra.BoundingBox3D(ub, lb)
1389 
1390  all_idxs = set(all_idxs) | collision_included_idxs
1391  all_idxs = all_idxs - collision_excluded_idxs
1392  debug = []
1393  print('shuffling', len(rigid_bodies), 'rigid bodies')
1394  for rb in rigid_bodies:
1395  if rb not in excluded_rigid_bodies:
1396  # gather particles to avoid with this transform
1397  if avoidcollision_rb:
1398  rb_idxs = set(rb.get_member_particle_indexes()) - \
1399  collision_excluded_idxs
1400  other_idxs = all_idxs - rb_idxs
1401 
1402  # iterate, trying to avoid collisions
1403  niter = 0
1404  while niter < niterations:
1405  rbxyz = (rb.get_x(), rb.get_y(), rb.get_z())
1406 
1407  # local transform
1408  if bounding_box:
1409  translation = IMP.algebra.get_random_vector_in(bb)
1410  # First move to origin
1411  transformation_orig = IMP.algebra.Transformation3D(
1413  -IMP.core.XYZ(rb).get_coordinates())
1414  IMP.core.transform(rb, transformation_orig)
1416  transformation = IMP.algebra.Transformation3D(rotation,
1417  translation)
1418 
1419  else:
1420  transformation = \
1422  rbxyz, max_translation, max_rotation)
1423 
1424  debug.append([rb, other_idxs if avoidcollision_rb else set()])
1425  IMP.core.transform(rb, transformation)
1426 
1427  # check collisions
1428  if avoidcollision_rb and other_idxs:
1429  mdl.update()
1430  npairs = len(gcpf.get_close_pairs(mdl,
1431  list(other_idxs),
1432  list(rb_idxs)))
1433  if npairs == 0:
1434  break
1435  else:
1436  niter += 1
1437  if verbose:
1438  print("shuffle_configuration: rigid body placed "
1439  "close to other %d particles, trying "
1440  "again..." % npairs)
1441  print("shuffle_configuration: rigid body name: "
1442  + rb.get_name())
1443  if niter == niterations:
1444  raise ValueError(
1445  "tried the maximum number of iterations to "
1446  "avoid collisions, increase the distance "
1447  "cutoff")
1448  else:
1449  break
1450 
1451  print('shuffling', len(flexible_beads), 'flexible beads')
1452  for fb in flexible_beads:
1453  # gather particles to avoid
1454  if avoidcollision_fb:
1455  fb_idxs = set(IMP.get_indexes([fb]))
1456  other_idxs = all_idxs - fb_idxs
1457  if not other_idxs:
1458  continue
1459 
1460  # iterate, trying to avoid collisions
1461  niter = 0
1462  while niter < niterations:
1463  if bounding_box:
1464  translation = IMP.algebra.get_random_vector_in(bb)
1465  transformation = IMP.algebra.Transformation3D(translation)
1466  else:
1467  fbxyz = IMP.core.XYZ(fb).get_coordinates()
1469  fbxyz, max_translation, max_rotation)
1470 
1471  # For gaussians, treat this fb as an rb
1473  memb = IMP.core.NonRigidMember(fb)
1474  xyz = memb.get_internal_coordinates()
1475  if bounding_box:
1476  # 'translation' is the new desired position in global
1477  # coordinates; we need to convert that to internal
1478  # coordinates first using the rigid body's ref frame
1479  rf = memb.get_rigid_body().get_reference_frame()
1480  glob_to_int = rf.get_transformation_from()
1481  memb.set_internal_coordinates(
1482  glob_to_int.get_transformed(translation))
1483  else:
1484  xyz_transformed = transformation.get_transformed(xyz)
1485  memb.set_internal_coordinates(xyz_transformed)
1486  debug.append([xyz, other_idxs if avoidcollision_fb else set()])
1487  else:
1488  d = IMP.core.XYZ(fb)
1489  if bounding_box:
1490  # Translate to origin first
1491  if IMP.core.RigidBody.get_is_setup(fb.get_particle()):
1493  IMP.core.RigidBody(fb.get_particle()),
1494  -d.get_coordinates())
1495  else:
1496  IMP.core.transform(d, -d.get_coordinates())
1497  d = IMP.core.XYZ(fb)
1498  debug.append([d, other_idxs if avoidcollision_fb else set()])
1499  if IMP.core.RigidBody.get_is_setup(fb.get_particle()):
1501  IMP.core.RigidBody(fb.get_particle()), transformation)
1502  else:
1503  IMP.core.transform(d, transformation)
1504 
1505  if avoidcollision_fb:
1506  mdl.update()
1507  npairs = len(gcpf.get_close_pairs(mdl,
1508  list(other_idxs),
1509  list(fb_idxs)))
1510 
1511  if npairs == 0:
1512  break
1513  else:
1514  niter += 1
1515  print("shuffle_configuration: floppy body placed close "
1516  "to other %d particles, trying again..." % npairs)
1517  if niter == niterations:
1518  raise ValueError(
1519  "tried the maximum number of iterations to avoid "
1520  "collisions, increase the distance cutoff")
1521  else:
1522  break
1523  if return_debug:
1524  return debug
1525 
1526 
1527 class ColorHierarchy(object):
1528 
1529  def __init__(self, hier):
1530  import matplotlib as mpl
1531  mpl.use('Agg')
1532  import matplotlib.pyplot as plt
1533  self.mpl = mpl
1534  self.plt = plt
1535 
1536  hier.ColorHierarchy = self
1537  self.hier = hier
1539  self.mols = [IMP.pmi.topology.PMIMoleculeHierarchy(mol)
1540  for mol in mols]
1541  self.method = self.nochange
1542  self.scheme = None
1543  self.first = None
1544  self.last = None
1545 
1546  def nochange(self):
1547  pass
1548 
1549  def get_color(self, fl):
1550  return IMP.display.Color(*self.scheme(fl)[0:3])
1551 
1552  def get_log_scale(self, fl):
1553  import math
1554  eps = 1.0
1555  return math.log(fl+eps)
1556 
1557  def color_by_resid(self):
1558  self.method = self.color_by_resid
1559  self.scheme = self.mpl.cm.rainbow
1560  for mol in self.mols:
1561  self.first = 1
1562  self.last = len(IMP.pmi.topology.PMIMoleculeHierarchy(
1563  mol).get_residue_indexes())
1564  for p in IMP.atom.get_leaves(mol):
1566  ri = IMP.atom.Residue(p).get_index()
1567  c = self.get_color(float(ri)/self.last)
1568  IMP.display.Colored(p).set_color(c)
1571  avr = sum(ris)/len(ris)
1572  c = self.get_color(float(avr)/self.last)
1573  IMP.display.Colored(p).set_color(c)
1574 
1575  def color_by_uncertainty(self):
1576  self.method = self.color_by_uncertainty
1577  self.scheme = self.mpl.cm.jet
1578  ps = IMP.atom.get_leaves(self.hier)
1579  unc_dict = {}
1580  for p in ps:
1582  u = IMP.pmi.Uncertainty(p).get_uncertainty()
1583  unc_dict[p] = u
1584  self.first = self.get_log_scale(1.0)
1585  self.last = self.get_log_scale(100.0)
1586  for p in unc_dict:
1587  value = self.get_log_scale(unc_dict[p])
1588  if value >= self.last:
1589  value = self.last
1590  if value <= self.first:
1591  value = self.first
1592  c = self.get_color((value-self.first) / (self.last-self.first))
1593  IMP.display.Colored(p).set_color(c)
1594 
1595  def get_color_bar(self, filename):
1596  import matplotlib as mpl
1597  mpl.use('Agg')
1598  import matplotlib.pyplot as plt
1599  plt.clf()
1600  fig = plt.figure(figsize=(8, 3))
1601  ax1 = fig.add_axes([0.05, 0.80, 0.9, 0.15])
1602 
1603  cmap = self.scheme
1604  norm = mpl.colors.Normalize(vmin=0.0, vmax=1.0)
1605 
1606  if self.method == self.color_by_uncertainty:
1607  angticks = [1.0, 2.5, 5.0, 10.0, 25.0, 50.0, 100.0]
1608  vvalues = []
1609  marks = []
1610  for at in angticks:
1611  vvalue = (self.get_log_scale(at)-self.first) \
1612  / (self.last-self.first)
1613  if vvalue <= 1.0 and vvalue >= 0.0:
1614  vvalues.append(vvalue)
1615  marks.append(str(at))
1616  cb1 = mpl.colorbar.ColorbarBase(
1617  ax1, cmap=cmap, norm=norm, ticks=vvalues,
1618  orientation='horizontal')
1619  print(self.first, self.last, marks, vvalues)
1620  cb1.ax.set_xticklabels(marks)
1621  cb1.set_label('Angstorm')
1622  plt.savefig(filename, dpi=150, transparent=True)
1623  plt.show()
1624 
1625 
1626 def color2rgb(colorname):
1627  """Given a Chimera color name or hex color value, return RGB"""
1628  d = {'aquamarine': (0.4980392156862745, 1.0, 0.8313725490196079),
1629  'black': (0.0, 0.0, 0.0),
1630  'blue': (0.0, 0.0, 1.0),
1631  'brown': (0.6470588235, 0.16470588235294117, 0.16470588235294117),
1632  'chartreuse': (0.4980392156862745, 1.0, 0.0),
1633  'coral': (1.0, 0.4980392156862745, 0.3137254901960784),
1634  'cornflower blue': (0.39215686, 0.58431372549, 0.9294117647058824),
1635  'cyan': (0.0, 1.0, 1.0),
1636  'dark cyan': (0.0, 0.5450980392156862, 0.5450980392156862),
1637  'dark gray': (0.6627450980, 0.6627450980392157, 0.6627450980392157),
1638  'dark green': (0.0, 0.39215686274509803, 0.0),
1639  'dark khaki': (0.74117647, 0.7176470588235294, 0.4196078431372549),
1640  'dark magenta': (0.5450980392156862, 0.0, 0.5450980392156862),
1641  'dark olive green': (0.333333333, 0.419607843, 0.1843137254901961),
1642  'dark red': (0.5450980392156862, 0.0, 0.0),
1643  'dark slate blue': (0.28235294, 0.239215686, 0.5450980392156862),
1644  'dark slate gray': (0.1843137, 0.30980392, 0.30980392156862746),
1645  'deep pink': (1.0, 0.0784313725490196, 0.5764705882352941),
1646  'deep sky blue': (0.0, 0.7490196078431373, 1.0),
1647  'dim gray': (0.41176470, 0.4117647058823529, 0.4117647058823529),
1648  'dodger blue': (0.11764705882352941, 0.5647058823529412, 1.0),
1649  'firebrick': (0.6980392, 0.13333333333333333, 0.13333333333333333),
1650  'forest green': (0.13333333, 0.5450980392156862, 0.13333333333333333),
1651  'gold': (1.0, 0.8431372549019608, 0.0),
1652  'goldenrod': (0.85490196, 0.6470588235294118, 0.12549019607843137),
1653  'gray': (0.7450980392156863, 0.7450980392156863, 0.7450980392156863),
1654  'green': (0.0, 1.0, 0.0),
1655  'hot pink': (1.0, 0.4117647058823529, 0.7058823529411765),
1656  'khaki': (0.9411764705882353, 0.9019607843137255, 0.5490196078431373),
1657  'light blue': (0.67843137, 0.8470588235294118, 0.9019607843137255),
1658  'light gray': (0.82745098, 0.8274509803921568, 0.8274509803921568),
1659  'light green': (0.56470588, 0.9333333333333333, 0.5647058823529412),
1660  'light sea green': (0.125490, 0.6980392156862745, 0.6666666666666666),
1661  'lime green': (0.1960784, 0.803921568627451, 0.19607843137254902),
1662  'magenta': (1.0, 0.0, 1.0),
1663  'medium blue': (0.1960784, 0.19607843137254902, 0.803921568627451),
1664  'medium purple': (0.57647, 0.4392156862745098, 0.8588235294117647),
1665  'navy blue': (0.0, 0.0, 0.5019607843137255),
1666  'olive drab': (0.4196078, 0.5568627450980392, 0.13725490196078433),
1667  'orange red': (1.0, 0.27058823529411763, 0.0),
1668  'orange': (1.0, 0.4980392156862745, 0.0),
1669  'orchid': (0.85490196, 0.4392156862745098, 0.8392156862745098),
1670  'pink': (1.0, 0.7529411764705882, 0.796078431372549),
1671  'plum': (0.8666666666666667, 0.6274509803921569, 0.8666666666666667),
1672  'purple': (0.62745098, 0.12549019607843137, 0.9411764705882353),
1673  'red': (1.0, 0.0, 0.0),
1674  'rosy brown': (0.7372549, 0.5607843137254902, 0.5607843137254902),
1675  'salmon': (0.980392, 0.5019607843137255, 0.4470588235294118),
1676  'sandy brown': (0.956862745, 0.6431372549019608, 0.3764705882352941),
1677  'sea green': (0.18039, 0.5450980392156862, 0.3411764705882353),
1678  'sienna': (0.6274509, 0.3215686274509804, 0.17647058823529413),
1679  'sky blue': (0.52941176, 0.807843137254902, 0.9215686274509803),
1680  'slate gray': (0.439215686, 0.50196078, 0.5647058823529412),
1681  'spring green': (0.0, 1.0, 0.4980392156862745),
1682  'steel blue': (0.2745098, 0.50980392, 0.70588235),
1683  'tan': (0.8235294117647058, 0.7058823529411765, 0.5490196078431373),
1684  'turquoise': (0.25098039, 0.87843137, 0.81568627),
1685  'violet red': (0.81568627, 0.125490196, 0.56470588235),
1686  'white': (1.0, 1.0, 1.0),
1687  'yellow': (1.0, 1.0, 0.0)}
1688  if colorname.startswith('#'):
1689  return tuple(int(colorname[i:i+2], 16) / 255. for i in (1, 3, 5))
1690  else:
1691  return d[colorname]
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: Molecule.h:35
def list_chunks_iterator
Yield successive length-sized chunks from a list.
Definition: tools.py:659
Simple 3D transformation class.
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: Residue.h:156
def select_at_all_resolutions
Perform selection using the usual keywords but return ALL resolutions (BEADS and GAUSSIANS).
Definition: tools.py:1134
Represent an RGB color.
Definition: Color.h:24
A decorator to associate a particle with a part of a protein/DNA/RNA.
Definition: Fragment.h:20
Extends the functionality of IMP.atom.Molecule.
def display_bonds
Decorate the sequence-consecutive particles from a PMI2 molecule with a bond, so that they appear con...
Definition: tools.py:1105
def add_script_provenance
Tag the given particle with the current Python script.
def get_restraint_set
Get a RestraintSet containing all PMI restraints added to the model.
Definition: tools.py:107
def remove
index can be a integer
Definition: tools.py:742
def shuffle_configuration
Shuffle particles.
Definition: tools.py:1314
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: atom/Atom.h:241
def get_particles_within_zone
Utility to retrieve particles from a hierarchy within a zone around a set of ps.
Definition: tools.py:1172
Set of Python classes to create a multi-state, multi-resolution IMP hierarchy.
def __init__
Constructor.
Definition: tools.py:127
static Weight setup_particle(Model *m, ParticleIndex pi)
Set up an empty Weight.
Definition: Weight.h:48
A decorator for a particle which has bonds.
Rotation3D get_random_rotation_3d(const Rotation3D &center, double distance)
Pick a rotation at random near the provided one.
def get_molecules
This function returns the parent molecule hierarchies of given objects.
Definition: tools.py:1231
An exception for an invalid usage of IMP.
Definition: exception.h:123
std::string get_module_version()
Return the version of this module, as a string.
Change color code to hexadecimal to rgb.
Definition: tools.py:805
This class stores integers in ordered compact lists eg: [[1,2,3],[6,7,8]] the methods help splitting ...
Definition: tools.py:679
The type of an atom.
static Surface setup_particle(Model *m, ParticleIndex pi)
Definition: Surface.h:45
Add uncertainty to a particle.
Definition: Uncertainty.h:24
void add_particle(RMF::FileHandle fh, Particle *hs)
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:731
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: XYZ.h:49
GenericHierarchies get_leaves(Hierarchy mhd)
Get all the leaves of the bit of hierarchy.
def get_prot_name_from_particle
Return the component name provided a particle and a list of names.
Definition: tools.py:545
Represent the root node of the global IMP.atom.Hierarchy.
def get_residue_gaps_in_hierarchy
Return the residue index gaps and contiguous segments in the hierarchy.
Definition: tools.py:419
def add_imp_provenance
Tag the given particle as being created by the current version of IMP.
def add
index can be a integer or a list of integers
Definition: tools.py:707
Vector3D get_random_vector_in(const Cylinder3D &c)
Generate a random vector in a cylinder with uniform density.
def get_all_leaves
Just get the leaves from a list of hierarchies.
Definition: tools.py:1127
Bond create_bond(Bonded a, Bonded b, Bond o)
Connect the two wrapped particles by a custom bond.
Object used to hold a set of restraints.
Definition: RestraintSet.h:37
Stores a named protein chain.
def color2rgb
Given a Chimera color name or hex color value, return RGB.
Definition: tools.py:1626
def input_adaptor
Adapt things for PMI (degrees of freedom, restraints, ...) Returns list of list of hierarchies...
Definition: tools.py:944
static bool get_is_setup(Model *m, ParticleIndex pi)
Definition: Fragment.h:46
def add_software_provenance
Tag the given particle with the software used to create it.
A decorator for keeping track of copies of a molecule.
Definition: Copy.h:28
ParticleIndexPairs get_indexes(const ParticlePairsTemp &ps)
Get the indexes from a list of particle pairs.
The standard decorator for manipulating molecular structures.
Ints get_index(const ParticlesTemp &particles, const Subset &subset, const Subsets &excluded)
def select_by_tuple_2
New tuple format: molname OR (start,stop,molname,copynum,statenum) Copy and state are optional...
Definition: tools.py:496
A decorator for a particle representing an atom.
Definition: atom/Atom.h:234
def scatter_and_gather
Synchronize data over a parallel run.
Definition: tools.py:608
void transform(XYZ a, const algebra::Transformation3D &tr)
Apply a transformation to the particle.
def add_restraint_to_model
Add a PMI restraint to the model.
Definition: tools.py:87
static Bonded setup_particle(Model *m, ParticleIndex pi)
def __init__
index can be a integer or a list of integers
Definition: tools.py:696
void load_frame(RMF::FileConstHandle file, RMF::FrameID frame)
Load the given RMF frame into the state of the linked objects.
A decorator for a particle with x,y,z coordinates.
Definition: XYZ.h:30
def get_flatten
Returns a flatten list.
Definition: tools.py:760
static Scale setup_particle(Model *m, ParticleIndex pi)
Definition: Scale.h:30
A base class for Keys.
Definition: Key.h:44
Collect timing information.
Definition: tools.py:122
A decorator for a particle that is part of a rigid body but not rigid.
Definition: rigid_bodies.h:750
def get_closest_residue_position
this function works with plain hierarchies, as read from the pdb, no multi-scale hierarchies ...
Definition: tools.py:381
Find all nearby pairs by testing all pairs.
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:163
A decorator for a residue.
Definition: Residue.h:135
DensityGrid get_grid(IMP::em::DensityMap *in)
Return a dense grid containing the voxels of the passed density map.
General purpose algebraic and geometric methods that are expected to be used by a wide variety of IMP...
static bool get_is_setup(const IMP::ParticleAdaptor &p)
static bool get_is_setup(Model *m, ParticleIndex p)
Check if the particle has the needed attributes for a cast to succeed.
The general base class for IMP exceptions.
Definition: exception.h:49
VectorD< 3 > Vector3D
Definition: VectorD.h:421
Rotation3D get_identity_rotation_3d()
Return a rotation that does not do anything.
Definition: Rotation3D.h:356
def get_densities
Given a list of PMI objects, returns all density hierarchies within these objects.
Definition: tools.py:1278
void link_hierarchies(RMF::FileConstHandle fh, const atom::Hierarchies &hs)
Class to handle individual particles of a Model object.
Definition: Particle.h:41
Bond get_bond(Bonded a, Bonded b)
Get the bond between two particles.
Stores a list of Molecules all with the same State index.
std::string get_data_path(std::string file_name)
Return the full path to one of this module's data files.
def set_coordinates_from_rmf
Extract frame from RMF file and fill coordinates.
Definition: tools.py:918
int get_copy_index(Hierarchy h)
Walk up the hierarchy to find the current copy index.
Python classes to represent, score, sample and analyze models.
static bool get_is_setup(Model *m, ParticleIndex pi)
Definition: Uncertainty.h:30
A decorator for a rigid body.
Definition: rigid_bodies.h:82
def cross_link_db_filter_parser
example '"{ID_Score}" > 28 AND "{Sample}" == "%10_1%" OR ":Sample}" == "%10_2%" OR ":Sample...
Definition: tools.py:327
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: core/Gaussian.h:47
def get_sorted_segments
Returns sequence-sorted segments array, each containing the first particle the last particle and the ...
Definition: tools.py:1079
def get_rbs_and_beads
Returns unique objects in original order.
Definition: tools.py:1207
Hierarchies get_leaves(const Selection &h)
A decorator for a molecule.
Definition: Molecule.h:24
Select hierarchy particles identified by the biological name.
Definition: Selection.h:66
Support for the RMF file format for storing hierarchical molecular data and markup.
def get_residue_indexes
Retrieve the residue indexes for the given particle.
Definition: tools.py:565
static bool get_is_setup(const IMP::ParticleAdaptor &p)
Definition: rigid_bodies.h:752
Warning for probably incorrect input parameters.
Transformation3D get_random_local_transformation(Vector3D origin, double max_translation=5., double max_angle_in_rad=0.26)
Get a local transformation.
Temporarily stores residue information, even without structure available.
Store objects in order they were added, but with default type.
Definition: tools.py:888
Inferential scoring building on methods developed as part of the Inferential Structure Determination ...
def sublist_iterator
Yield all sublists of length >= lmin and <= lmax.
Definition: tools.py:640
A particle with a color.
Definition: Colored.h:23