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