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