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