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