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