IMP logo
IMP Reference Guide  2.17.0
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=1, lmax=None):
645  '''
646  Yield all sublists of length >= lmin and <= lmax
647  '''
648  if lmax is None:
649  lmax = len(ls)
650  n = len(ls)
651  for i in range(n):
652  for j in range(i + lmin, min(n + 1, i + 1 + lmax)):
653  yield ls[i:j]
654 
655 
656 def flatten_list(ls):
657  return [item for sublist in ls for item in sublist]
658 
659 
660 def list_chunks_iterator(list, length):
661  """ Yield successive length-sized chunks from a list.
662  """
663  for i in range(0, len(list), length):
664  yield list[i:i + length]
665 
666 
667 def chunk_list_into_segments(seq, num):
668  seq = list(seq)
669  avg = len(seq) / float(num)
670  out = []
671  last = 0.0
672 
673  while last < len(seq):
674  out.append(seq[int(last):int(last + avg)])
675  last += avg
676 
677  return out
678 
679 
680 class Segments(object):
681 
682  ''' This class stores integers
683  in ordered compact lists eg:
684  [[1,2,3],[6,7,8]]
685  the methods help splitting and merging the internal lists
686  Example:
687  s=Segments([1,2,3]) is [[1,2,3]]
688  s.add(4) is [[1,2,3,4]] (add right)
689  s.add(3) is [[1,2,3,4]] (item already existing)
690  s.add(7) is [[1,2,3,4],[7]] (new list)
691  s.add([8,9]) is [[1,2,3,4],[7,8,9]] (add item right)
692  s.add([5,6]) is [[1,2,3,4,5,6,7,8,9]] (merge)
693  s.remove(3) is [[1,2],[4,5,6,7,8,9]] (split)
694  etc.
695  '''
696 
697  def __init__(self, index):
698  '''index can be a integer or a list of integers '''
699  if isinstance(index, int):
700  self.segs = [[index]]
701  elif isinstance(index, list):
702  self.segs = [[index[0]]]
703  for i in index[1:]:
704  self.add(i)
705  else:
706  raise TypeError("index must be an int or list of ints")
707 
708  def add(self, index):
709  '''index can be a integer or a list of integers '''
710  if isinstance(index, (int, numpy.int32, numpy.int64)):
711  mergeleft = None
712  mergeright = None
713  for n, s in enumerate(self.segs):
714  if index in s:
715  return 0
716  else:
717  if s[0]-index == 1:
718  mergeleft = n
719  if index-s[-1] == 1:
720  mergeright = n
721  if mergeright is None and mergeleft is None:
722  self.segs.append([index])
723  if mergeright is not None and mergeleft is None:
724  self.segs[mergeright].append(index)
725  if mergeleft is not None and mergeright is None:
726  self.segs[mergeleft] = [index]+self.segs[mergeleft]
727  if mergeleft is not None and mergeright is not None:
728  self.segs[mergeright] = \
729  self.segs[mergeright]+[index]+self.segs[mergeleft]
730  del self.segs[mergeleft]
731 
732  for n in range(len(self.segs)):
733  self.segs[n].sort()
734 
735  self.segs.sort(key=lambda tup: tup[0])
736 
737  elif isinstance(index, list):
738  for i in index:
739  self.add(i)
740  else:
741  raise TypeError("index must be an int or list of ints")
742 
743  def remove(self, index):
744  '''index can be a integer'''
745  for n, s in enumerate(self.segs):
746  if index in s:
747  if s[0] == index:
748  self.segs[n] = s[1:]
749  elif s[-1] == index:
750  self.segs[n] = s[:-1]
751  else:
752  i = self.segs[n].index(index)
753  self.segs[n] = s[:i]
754  self.segs.append(s[i+1:])
755  for n in range(len(self.segs)):
756  self.segs[n].sort()
757  if len(self.segs[n]) == 0:
758  del self.segs[n]
759  self.segs.sort(key=lambda tup: tup[0])
760 
761  def get_flatten(self):
762  ''' Returns a flatten list '''
763  return [item for sublist in self.segs for item in sublist]
764 
765  def __repr__(self):
766  ret_tmp = "["
767  for seg in self.segs:
768  ret_tmp += str(seg[0])+"-"+str(seg[-1])+","
769  ret = ret_tmp[:-1]+"]"
770  return ret
771 
772 #
773 # Tools to simulate data
774 #
775 
776 
777 def normal_density_function(expected_value, sigma, x):
778  return (
779  1 / math.sqrt(2 * math.pi) / sigma *
780  math.exp(-(x - expected_value) ** 2 / 2 / sigma / sigma)
781  )
782 
783 
784 def log_normal_density_function(expected_value, sigma, x):
785  return (
786  1 / math.sqrt(2 * math.pi) / sigma / x *
787  math.exp(-(math.log(x / expected_value) ** 2 / 2 / sigma / sigma))
788  )
789 
790 
791 def print_multicolumn(list_of_strings, ncolumns=2, truncate=40):
792 
793  ls = list_of_strings
794 
795  cols = ncolumns
796  # add empty entries after ls
797  for i in range(len(ls) % cols):
798  ls.append(" ")
799 
800  split = [ls[i:i + len(ls) // cols]
801  for i in range(0, len(ls), len(ls) // cols)]
802  for row in zip(*split):
803  print("".join(str.ljust(i, truncate) for i in row))
804 
805 
806 class ColorChange(object):
807  '''Change color code to hexadecimal to rgb'''
808  def __init__(self):
809  self._NUMERALS = '0123456789abcdefABCDEF'
810  self._HEXDEC = dict((v, int(v, 16)) for v in
811  (x+y for x in self._NUMERALS
812  for y in self._NUMERALS))
813  self.LOWERCASE, self.UPPERCASE = 'x', 'X'
814 
815  def rgb(self, triplet):
816  return (float(self._HEXDEC[triplet[0:2]]),
817  float(self._HEXDEC[triplet[2:4]]),
818  float(self._HEXDEC[triplet[4:6]]))
819 
820  def triplet(self, rgb, lettercase=None):
821  if lettercase is None:
822  lettercase = self.LOWERCASE
823  return format(rgb[0] << 16 | rgb[1] << 8 | rgb[2], '06'+lettercase)
824 
825 
826 # -------------- Collections --------------- #
827 class OrderedSet(MutableSet):
828 
829  def __init__(self, iterable=None):
830  self.end = end = []
831  end += [None, end, end] # sentinel node for doubly linked list
832  self.map = {} # key --> [key, prev, next]
833  if iterable is not None:
834  self |= iterable
835 
836  def __len__(self):
837  return len(self.map)
838 
839  def __contains__(self, key):
840  return key in self.map
841 
842  def add(self, key):
843  if key not in self.map:
844  end = self.end
845  curr = end[1]
846  curr[2] = end[1] = self.map[key] = [key, curr, end]
847 
848  def discard(self, key):
849  if key in self.map:
850  key, prev, next = self.map.pop(key)
851  prev[2] = next
852  next[1] = prev
853 
854  def __iter__(self):
855  end = self.end
856  curr = end[2]
857  while curr is not end:
858  yield curr[0]
859  curr = curr[2]
860 
861  def __reversed__(self):
862  end = self.end
863  curr = end[1]
864  while curr is not end:
865  yield curr[0]
866  curr = curr[1]
867 
868  def pop(self, last=True):
869  if not self:
870  raise KeyError('set is empty')
871  if last:
872  key = self.end[1][0]
873  else:
874  key = self.end[2][0]
875  self.discard(key)
876  return key
877 
878  def __repr__(self):
879  if not self:
880  return '%s()' % (self.__class__.__name__,)
881  return '%s(%r)' % (self.__class__.__name__, list(self))
882 
883  def __eq__(self, other):
884  if isinstance(other, OrderedSet):
885  return len(self) == len(other) and list(self) == list(other)
886  return set(self) == set(other)
887 
888 
889 class OrderedDefaultDict(OrderedDict):
890  """Store objects in order they were added, but with default type.
891  Source: http://stackoverflow.com/a/4127426/2608793
892  """
893  def __init__(self, *args, **kwargs):
894  if not args:
895  self.default_factory = None
896  else:
897  if not (args[0] is None or callable(args[0])):
898  raise TypeError('first argument must be callable or None')
899  self.default_factory = args[0]
900  args = args[1:]
901  super(OrderedDefaultDict, self).__init__(*args, **kwargs)
902 
903  def __missing__(self, key):
904  if self.default_factory is None:
905  raise KeyError(key)
906  self[key] = default = self.default_factory()
907  return default
908 
909  def __reduce__(self): # optional, for pickle support
910  args = (self.default_factory,) if self.default_factory else ()
911  if sys.version_info[0] >= 3:
912  return self.__class__, args, None, None, self.items()
913  else:
914  return self.__class__, args, None, None, self.iteritems()
915 
916 
917 # -------------- PMI2 Tools --------------- #
918 
919 def set_coordinates_from_rmf(hier, rmf_fn, frame_num=0):
920  """Extract frame from RMF file and fill coordinates. Must be identical
921  topology.
922 
923  @param hier The (System) hierarchy to fill (e.g. after you've built it)
924  @param rmf_fn The file to extract from
925  @param frame_num The frame number to extract
926  """
927  rh = RMF.open_rmf_file_read_only(rmf_fn)
928  IMP.rmf.link_hierarchies(rh, [hier])
929  IMP.rmf.load_frame(rh, RMF.FrameID(frame_num))
930  del rh
931 
932 
933 def input_adaptor(stuff, pmi_resolution=0, flatten=False, selection_tuple=None,
934  warn_about_slices=True):
935  """Adapt things for PMI (degrees of freedom, restraints, ...)
936  Returns list of list of hierarchies, separated into Molecules if possible.
937  The input can be a list, or a list of lists (iterable of ^1 or
938  iterable of ^2)
939  (iterable of ^2) Hierarchy -> returns input as list of list of hierarchies,
940  only one entry, not grouped by molecules.
941  (iterable of ^2) PMI::System/State/Molecule/TempResidue ->
942  returns residue hierarchies, grouped in molecules, at requested
943  resolution
944 
945  @param stuff Can be one of the following inputs:
946  IMP Hierarchy, PMI System/State/Molecule/TempResidue, or a
947  list/set (of list/set) of them.
948  Must be uniform input, however. No mixing object types.
949  @param pmi_resolution For selecting, only does it if you pass PMI
950  objects. Set it to "all" if you want all resolutions!
951  @param flatten Set to True if you just want all hierarchies in one list.
952  @param warn_about_slices Print a warning if you are requesting only part
953  of a bead. Sometimes you just don't care!
954  @note since this relies on IMP::atom::Selection, this will not return
955  any objects if they weren't built! But there should be no problem
956  if you request unbuilt residues - they should be ignored.
957  """
958 
959  if stuff is None:
960  return stuff
961 
962  if hasattr(stuff, '__iter__'):
963  if len(stuff) == 0:
964  return stuff
965  thelist = list(stuff)
966 
967  # iter of iter of should be ok
968  if all(hasattr(el, '__iter__') for el in thelist):
969  thelist = [i for sublist in thelist for i in sublist]
970  elif any(hasattr(el, '__iter__') for el in thelist):
971  raise Exception('input_adaptor: input_object must be a list '
972  'or a list of lists')
973 
974  stuff = thelist
975  else:
976  stuff = [stuff]
977 
978  # check that it is a hierarchy homogenously:
979  try:
980  is_hierarchy = all(IMP.atom.Hierarchy.get_is_setup(s) for s in stuff)
981  except (NotImplementedError, TypeError):
982  is_hierarchy = False
983  # get the other types homogenously
984  is_system = all(isinstance(s, IMP.pmi.topology.System) for s in stuff)
985  is_state = all(isinstance(s, IMP.pmi.topology.State) for s in stuff)
986  is_molecule = all(isinstance(s, IMP.pmi.topology.Molecule) for s in stuff)
987  is_temp_residue = all(isinstance(s, IMP.pmi.topology.TempResidue)
988  for s in stuff)
989 
990  # now that things are ok, do selection if requested
991  hier_list = []
992  pmi_input = False
993  if is_system or is_state or is_molecule or is_temp_residue:
994  # if PMI, perform selection using gathered indexes
995  pmi_input = True
996  # key is Molecule object, value are residues
997  indexes_per_mol = OrderedDefaultDict(list)
998  if is_system:
999  for system in stuff:
1000  for state in system.get_states():
1001  mdict = state.get_molecules()
1002  for molname in mdict:
1003  for copy in mdict[molname]:
1004  indexes_per_mol[copy] += \
1005  [r.get_index() for r in copy.get_residues()]
1006  elif is_state:
1007  for state in stuff:
1008  mdict = state.get_molecules()
1009  for molname in mdict:
1010  for copy in mdict[molname]:
1011  indexes_per_mol[copy] += [r.get_index()
1012  for r in copy.get_residues()]
1013  elif is_molecule:
1014  for molecule in stuff:
1015  indexes_per_mol[molecule] += [r.get_index()
1016  for r in molecule.get_residues()]
1017  else: # is_temp_residue
1018  for tempres in stuff:
1019  indexes_per_mol[tempres.get_molecule()].append(
1020  tempres.get_index())
1021  for mol in indexes_per_mol:
1022  if pmi_resolution == 'all':
1023  # because you select from the molecule,
1024  # this will start the search from the base resolution
1026  mol.get_hierarchy(), residue_indexes=indexes_per_mol[mol])
1027  else:
1028  sel = IMP.atom.Selection(mol.get_hierarchy(),
1029  resolution=pmi_resolution,
1030  residue_indexes=indexes_per_mol[mol])
1031  ps = sel.get_selected_particles()
1032 
1033  # check that you don't have any incomplete fragments!
1034  if warn_about_slices:
1035  rset = set(indexes_per_mol[mol])
1036  for p in ps:
1038  fset = set(IMP.atom.Fragment(p).get_residue_indexes())
1039  if not fset <= rset:
1040  minset = min(fset)
1041  maxset = max(fset)
1042  found = fset & rset
1043  minf = min(found)
1044  maxf = max(found)
1045  resbreak = maxf if minf == minset else minset-1
1046  warnings.warn(
1047  'You are trying to select only part of the '
1048  'bead %s:%i-%i. The residues you requested '
1049  'are %i-%i. You can fix this by: '
1050  '1) requesting the whole bead/none of it; or'
1051  '2) break the bead up by passing '
1052  'bead_extra_breaks=[\'%i\'] in '
1053  'molecule.add_representation()'
1054  % (mol.get_name(), minset, maxset, minf, maxf,
1055  resbreak), IMP.pmi.ParameterWarning)
1056  hier_list.append([IMP.atom.Hierarchy(p) for p in ps])
1057  elif is_hierarchy:
1058  # check
1059  ps = []
1060  if pmi_resolution == 'all':
1061  for h in stuff:
1062  ps += select_at_all_resolutions(h)
1063  else:
1064  for h in stuff:
1065  ps += IMP.atom.Selection(
1066  h, resolution=pmi_resolution).get_selected_particles()
1067  hier_list = [IMP.atom.Hierarchy(p) for p in ps]
1068  if not flatten:
1069  hier_list = [hier_list]
1070  else:
1071  raise Exception('input_adaptor: you passed something of wrong type '
1072  'or a list with mixed types')
1073 
1074  if flatten and pmi_input:
1075  return [h for sublist in hier_list for h in sublist]
1076  else:
1077  return hier_list
1078 
1079 
1081  """Returns sequence-sorted segments array, each containing the first
1082  particle the last particle and the first residue index."""
1083 
1084  from operator import itemgetter
1085  hiers = IMP.pmi.tools.input_adaptor(mol)
1086  if len(hiers) > 1:
1087  raise ValueError("only pass stuff from one Molecule, please")
1088  hiers = hiers[0]
1089  segs = []
1090  for h in hiers:
1091  try:
1092  start = IMP.atom.Hierarchy(h).get_children()[0]
1093  except IndexError:
1094  start = IMP.atom.Hierarchy(h)
1095 
1096  try:
1097  end = IMP.atom.Hierarchy(h).get_children()[-1]
1098  except IndexError:
1099  end = IMP.atom.Hierarchy(h)
1100 
1101  startres = IMP.pmi.tools.get_residue_indexes(start)[0]
1102  segs.append((start, end, startres))
1103  return sorted(segs, key=itemgetter(2))
1104 
1105 
1106 def display_bonds(mol):
1107  """Decorate the sequence-consecutive particles from a PMI2 molecule
1108  with a bond, so that they appear connected in the rmf file"""
1109  SortedSegments = get_sorted_segments(mol)
1110  for x in range(len(SortedSegments) - 1):
1111 
1112  last = SortedSegments[x][1]
1113  first = SortedSegments[x + 1][0]
1114 
1115  p1 = last.get_particle()
1116  p2 = first.get_particle()
1117  if not IMP.atom.Bonded.get_is_setup(p1):
1119  if not IMP.atom.Bonded.get_is_setup(p2):
1121 
1124  IMP.atom.Bonded(p1),
1125  IMP.atom.Bonded(p2), 1)
1126 
1127 
1128 def get_all_leaves(list_of_hs):
1129  """ Just get the leaves from a list of hierarchies """
1130  lvs = list(itertools.chain.from_iterable(
1131  IMP.atom.get_leaves(item) for item in list_of_hs))
1132  return lvs
1133 
1134 
1135 def select_at_all_resolutions(hier=None, hiers=None, **kwargs):
1136  """Perform selection using the usual keywords but return ALL
1137  resolutions (BEADS and GAUSSIANS).
1138  Returns in flat list!
1139  """
1140 
1141  if hiers is None:
1142  hiers = []
1143  if hier is not None:
1144  hiers.append(hier)
1145  if len(hiers) == 0:
1146  warnings.warn("You passed nothing to select_at_all_resolutions()",
1148  return []
1149  ret = OrderedSet()
1150  for hsel in hiers:
1151  try:
1152  htest = IMP.atom.Hierarchy.get_is_setup(hsel)
1153  except: # noqa: E722
1154  raise Exception('select_at_all_resolutions: you have to pass '
1155  'an IMP Hierarchy')
1156  if not htest:
1157  raise Exception('select_at_all_resolutions: you have to pass '
1158  'an IMP Hierarchy')
1159  if 'resolution' in kwargs or 'representation_type' in kwargs:
1160  raise Exception("don't pass resolution or representation_type "
1161  "to this function")
1162  selB = IMP.atom.Selection(hsel, resolution=IMP.atom.ALL_RESOLUTIONS,
1163  representation_type=IMP.atom.BALLS,
1164  **kwargs)
1165  selD = IMP.atom.Selection(hsel, resolution=IMP.atom.ALL_RESOLUTIONS,
1166  representation_type=IMP.atom.DENSITIES,
1167  **kwargs)
1168  ret |= OrderedSet(selB.get_selected_particles())
1169  ret |= OrderedSet(selD.get_selected_particles())
1170  return list(ret)
1171 
1172 
1174  target_ps,
1175  sel_zone,
1176  entire_residues,
1177  exclude_backbone):
1178  """Utility to retrieve particles from a hierarchy within a
1179  zone around a set of ps.
1180  @param hier The hierarchy in which to look for neighbors
1181  @param target_ps The particles for zoning
1182  @param sel_zone The maximum distance
1183  @param entire_residues If True, will grab entire residues
1184  @param exclude_backbone If True, will only return sidechain particles
1185  """
1186 
1187  test_sel = IMP.atom.Selection(hier)
1188  backbone_types = ['C', 'N', 'CB', 'O']
1189  if exclude_backbone:
1190  test_sel -= IMP.atom.Selection(
1191  hier, atom_types=[IMP.atom.AtomType(n) for n in backbone_types])
1192  test_ps = test_sel.get_selected_particles()
1193  nn = IMP.algebra.NearestNeighbor3D([IMP.core.XYZ(p).get_coordinates()
1194  for p in test_ps])
1195  zone = set()
1196  for target in target_ps:
1197  zone |= set(nn.get_in_ball(IMP.core.XYZ(target).get_coordinates(),
1198  sel_zone))
1199  zone_ps = [test_ps[z] for z in zone]
1200  if entire_residues:
1201  final_ps = set()
1202  for z in zone_ps:
1203  final_ps |= set(IMP.atom.Hierarchy(z).get_parent().get_children())
1204  zone_ps = [h.get_particle() for h in final_ps]
1205  return zone_ps
1206 
1207 
1209  """Returns unique objects in original order"""
1210  rbs = set()
1211  beads = []
1212  rbs_ordered = []
1213  if not hasattr(hiers, '__iter__'):
1214  hiers = [hiers]
1215  for p in get_all_leaves(hiers):
1217  rb = IMP.core.RigidMember(p).get_rigid_body()
1218  if rb not in rbs:
1219  rbs.add(rb)
1220  rbs_ordered.append(rb)
1222  rb = IMP.core.NonRigidMember(p).get_rigid_body()
1223  if rb not in rbs:
1224  rbs.add(rb)
1225  rbs_ordered.append(rb)
1226  beads.append(p)
1227  else:
1228  beads.append(p)
1229  return rbs_ordered, beads
1230 
1231 
1232 def get_molecules(input_objects):
1233  "This function returns the parent molecule hierarchies of given objects"
1234  stuff = input_adaptor(input_objects, pmi_resolution='all', flatten=True)
1235  molecules = set()
1236  for h in stuff:
1237  is_root = False
1238  is_molecule = False
1239  while not (is_molecule or is_root):
1240  root = IMP.atom.get_root(h)
1241  if root == h:
1242  is_root = True
1243  is_molecule = IMP.atom.Molecule.get_is_setup(h)
1244  if is_molecule:
1245  molecules.add(IMP.atom.Molecule(h))
1246  h = h.get_parent()
1247  return list(molecules)
1248 
1249 
1250 def get_molecules_dictionary(input_objects):
1251  moldict = defaultdict(list)
1252  for mol in IMP.pmi.tools.get_molecules(input_objects):
1253  name = mol.get_name()
1254  moldict[name].append(mol)
1255 
1256  for mol in moldict:
1257  moldict[mol].sort(key=lambda x: IMP.atom.Copy(x).get_copy_index())
1258  return moldict
1259 
1260 
1261 def get_molecules_dictionary_by_copy(input_objects):
1262  moldict = defaultdict(dict)
1263  for mol in IMP.pmi.tools.get_molecules(input_objects):
1264  name = mol.get_name()
1265  c = IMP.atom.Copy(mol).get_copy_index()
1266  moldict[name][c] = mol
1267  return moldict
1268 
1269 
1270 def get_selections_dictionary(input_objects):
1271  moldict = IMP.pmi.tools.get_molecules_dictionary(input_objects)
1272  seldict = defaultdict(list)
1273  for name, mols in moldict.items():
1274  for m in mols:
1275  seldict[name].append(IMP.atom.Selection(m))
1276  return seldict
1277 
1278 
1279 def get_densities(input_objects):
1280  """Given a list of PMI objects, returns all density hierarchies within
1281  these objects. The output of this function can be inputted into
1282  things such as EM restraints. This function is intended to gather
1283  density particles appended to molecules (and not other hierarchies
1284  which might have been appended to the root node directly).
1285  """
1286  # Note that Densities can only be selected at the Root or Molecule
1287  # level and not at the Leaves level.
1288  # we'll first get all molecule hierarchies corresponding to the leaves.
1289  molecules = get_molecules(input_objects)
1290  densities = []
1291  for i in molecules:
1292  densities += IMP.atom.Selection(
1293  i, representation_type=IMP.atom.DENSITIES).get_selected_particles()
1294  return densities
1295 
1296 
1297 def shuffle_configuration(objects,
1298  max_translation=300., max_rotation=2.0 * math.pi,
1299  avoidcollision_rb=True, avoidcollision_fb=False,
1300  cutoff=10.0, niterations=100,
1301  bounding_box=None,
1302  excluded_rigid_bodies=[],
1303  hierarchies_excluded_from_collision=[],
1304  hierarchies_included_in_collision=[],
1305  verbose=False,
1306  return_debug=False):
1307  """Shuffle particles. Used to restart the optimization.
1308  The configuration of the system is initialized by placing each
1309  rigid body and each bead randomly in a box. If `bounding_box` is
1310  specified, the particles are placed inside this box; otherwise, each
1311  particle is displaced by up to max_translation angstroms, and randomly
1312  rotated. Effort is made to place particles far enough from each other to
1313  prevent any steric clashes.
1314  @param objects Can be one of the following inputs:
1315  IMP Hierarchy, PMI System/State/Molecule/TempResidue, or
1316  a list/set of them
1317  @param max_translation Max translation (rbs and flexible beads)
1318  @param max_rotation Max rotation (rbs only)
1319  @param avoidcollision_rb check if the particle/rigid body was
1320  placed close to another particle; uses the optional
1321  arguments cutoff and niterations
1322  @param avoidcollision_fb Advanced. Generally you want this False because
1323  it's hard to shuffle beads.
1324  @param cutoff Distance less than this is a collision
1325  @param niterations How many times to try avoiding collision
1326  @param bounding_box Only shuffle particles within this box.
1327  Defined by ((x1,y1,z1),(x2,y2,z2)).
1328  @param excluded_rigid_bodies Don't shuffle these rigid body objects
1329  @param hierarchies_excluded_from_collision Don't count collision
1330  with these bodies
1331  @param hierarchies_included_in_collision Hierarchies that are not
1332  shuffled, but should be included in collision calculation
1333  (for fixed regions)
1334  @param verbose Give more output
1335  @note Best to only call this function after you've set up degrees
1336  of freedom
1337  For debugging purposes, returns: <shuffled indexes>,
1338  <collision avoided indexes>
1339  """
1340 
1341  # checking input
1342  hierarchies = IMP.pmi.tools.input_adaptor(objects,
1343  pmi_resolution='all',
1344  flatten=True)
1345  rigid_bodies, flexible_beads = get_rbs_and_beads(hierarchies)
1346  if len(rigid_bodies) > 0:
1347  mdl = rigid_bodies[0].get_model()
1348  elif len(flexible_beads) > 0:
1349  mdl = flexible_beads[0].get_model()
1350  else:
1351  raise Exception("Could not find any particles in the hierarchy")
1352  if len(rigid_bodies) == 0:
1353  print("shuffle_configuration: rigid bodies were not initialized")
1354 
1355  # gather all particles
1357  gcpf.set_distance(cutoff)
1358 
1359  # Add particles from excluded hierarchies to excluded list
1360  collision_excluded_hierarchies = IMP.pmi.tools.input_adaptor(
1361  hierarchies_excluded_from_collision, pmi_resolution='all',
1362  flatten=True)
1363 
1364  collision_included_hierarchies = IMP.pmi.tools.input_adaptor(
1365  hierarchies_included_in_collision, pmi_resolution='all', flatten=True)
1366 
1367  collision_excluded_idxs = set(
1368  leaf.get_particle().get_index()
1369  for h in collision_excluded_hierarchies
1370  for leaf in IMP.core.get_leaves(h))
1371 
1372  collision_included_idxs = set(
1373  leaf.get_particle().get_index()
1374  for h in collision_included_hierarchies
1375  for leaf in IMP.core.get_leaves(h))
1376 
1377  # Excluded collision with Gaussians
1378  all_idxs = [] # expand to representations?
1379  for p in IMP.pmi.tools.get_all_leaves(hierarchies):
1381  all_idxs.append(p.get_particle_index())
1383  collision_excluded_idxs.add(p.get_particle_index())
1384 
1385  if bounding_box is not None:
1386  ((x1, y1, z1), (x2, y2, z2)) = bounding_box
1387  ub = IMP.algebra.Vector3D(x1, y1, z1)
1388  lb = IMP.algebra.Vector3D(x2, y2, z2)
1389  bb = IMP.algebra.BoundingBox3D(ub, lb)
1390 
1391  all_idxs = set(all_idxs) | collision_included_idxs
1392  all_idxs = all_idxs - collision_excluded_idxs
1393  debug = []
1394  print('shuffling', len(rigid_bodies), 'rigid bodies')
1395  for rb in rigid_bodies:
1396  if rb not in excluded_rigid_bodies:
1397  # gather particles to avoid with this transform
1398  if avoidcollision_rb:
1399  rb_idxs = set(rb.get_member_particle_indexes()) - \
1400  collision_excluded_idxs
1401  other_idxs = all_idxs - rb_idxs
1402 
1403  # iterate, trying to avoid collisions
1404  niter = 0
1405  while niter < niterations:
1406  rbxyz = (rb.get_x(), rb.get_y(), rb.get_z())
1407 
1408  # local transform
1409  if bounding_box:
1410  translation = IMP.algebra.get_random_vector_in(bb)
1411  # First move to origin
1412  transformation_orig = IMP.algebra.Transformation3D(
1414  -IMP.core.XYZ(rb).get_coordinates())
1415  IMP.core.transform(rb, transformation_orig)
1417  transformation = IMP.algebra.Transformation3D(rotation,
1418  translation)
1419 
1420  else:
1421  transformation = \
1423  rbxyz, max_translation, max_rotation)
1424 
1425  debug.append([rb, other_idxs if avoidcollision_rb else set()])
1426  IMP.core.transform(rb, transformation)
1427 
1428  # check collisions
1429  if avoidcollision_rb and other_idxs:
1430  mdl.update()
1431  npairs = len(gcpf.get_close_pairs(mdl,
1432  list(other_idxs),
1433  list(rb_idxs)))
1434  if npairs == 0:
1435  break
1436  else:
1437  niter += 1
1438  if verbose:
1439  print("shuffle_configuration: rigid body placed "
1440  "close to other %d particles, trying "
1441  "again..." % npairs)
1442  print("shuffle_configuration: rigid body name: "
1443  + rb.get_name())
1444  if niter == niterations:
1445  raise ValueError(
1446  "tried the maximum number of iterations to "
1447  "avoid collisions, increase the distance "
1448  "cutoff")
1449  else:
1450  break
1451 
1452  print('shuffling', len(flexible_beads), 'flexible beads')
1453  for fb in flexible_beads:
1454  # gather particles to avoid
1455  if avoidcollision_fb:
1456  fb_idxs = set(IMP.get_indexes([fb]))
1457  other_idxs = all_idxs - fb_idxs
1458  if not other_idxs:
1459  continue
1460 
1461  # iterate, trying to avoid collisions
1462  niter = 0
1463  while niter < niterations:
1464  if bounding_box:
1465  translation = IMP.algebra.get_random_vector_in(bb)
1466  transformation = IMP.algebra.Transformation3D(translation)
1467  else:
1468  fbxyz = IMP.core.XYZ(fb).get_coordinates()
1470  fbxyz, max_translation, max_rotation)
1471 
1472  # For gaussians, treat this fb as an rb
1474  memb = IMP.core.NonRigidMember(fb)
1475  xyz = memb.get_internal_coordinates()
1476  if bounding_box:
1477  # 'translation' is the new desired position in global
1478  # coordinates; we need to convert that to internal
1479  # coordinates first using the rigid body's ref frame
1480  rf = memb.get_rigid_body().get_reference_frame()
1481  glob_to_int = rf.get_transformation_from()
1482  memb.set_internal_coordinates(
1483  glob_to_int.get_transformed(translation))
1484  else:
1485  xyz_transformed = transformation.get_transformed(xyz)
1486  memb.set_internal_coordinates(xyz_transformed)
1487  debug.append([xyz, other_idxs if avoidcollision_fb else set()])
1488  else:
1489  d = IMP.core.XYZ(fb)
1490  if bounding_box:
1491  # Translate to origin first
1492  if IMP.core.RigidBody.get_is_setup(fb.get_particle()):
1494  IMP.core.RigidBody(fb.get_particle()),
1495  -d.get_coordinates())
1496  else:
1497  IMP.core.transform(d, -d.get_coordinates())
1498  d = IMP.core.XYZ(fb)
1499  debug.append([d, other_idxs if avoidcollision_fb else set()])
1500  if IMP.core.RigidBody.get_is_setup(fb.get_particle()):
1502  IMP.core.RigidBody(fb.get_particle()), transformation)
1503  else:
1504  IMP.core.transform(d, transformation)
1505 
1506  if avoidcollision_fb:
1507  mdl.update()
1508  npairs = len(gcpf.get_close_pairs(mdl,
1509  list(other_idxs),
1510  list(fb_idxs)))
1511 
1512  if npairs == 0:
1513  break
1514  else:
1515  niter += 1
1516  print("shuffle_configuration: floppy body placed close "
1517  "to other %d particles, trying again..." % npairs)
1518  if niter == niterations:
1519  raise ValueError(
1520  "tried the maximum number of iterations to avoid "
1521  "collisions, increase the distance cutoff")
1522  else:
1523  break
1524  if return_debug:
1525  return debug
1526 
1527 
1528 class ColorHierarchy(object):
1529 
1530  def __init__(self, hier):
1531  import matplotlib as mpl
1532  mpl.use('Agg')
1533  import matplotlib.pyplot as plt
1534  self.mpl = mpl
1535  self.plt = plt
1536 
1537  hier.ColorHierarchy = self
1538  self.hier = hier
1540  self.mols = [IMP.pmi.topology.PMIMoleculeHierarchy(mol)
1541  for mol in mols]
1542  self.method = self.nochange
1543  self.scheme = None
1544  self.first = None
1545  self.last = None
1546 
1547  def nochange(self):
1548  pass
1549 
1550  def get_color(self, fl):
1551  return IMP.display.Color(*self.scheme(fl)[0:3])
1552 
1553  def get_log_scale(self, fl):
1554  import math
1555  eps = 1.0
1556  return math.log(fl+eps)
1557 
1558  def color_by_resid(self):
1559  self.method = self.color_by_resid
1560  self.scheme = self.mpl.cm.rainbow
1561  for mol in self.mols:
1562  self.first = 1
1563  self.last = len(IMP.pmi.topology.PMIMoleculeHierarchy(
1564  mol).get_residue_indexes())
1565  for p in IMP.atom.get_leaves(mol):
1567  ri = IMP.atom.Residue(p).get_index()
1568  c = self.get_color(float(ri)/self.last)
1569  IMP.display.Colored(p).set_color(c)
1572  avr = sum(ris)/len(ris)
1573  c = self.get_color(float(avr)/self.last)
1574  IMP.display.Colored(p).set_color(c)
1575 
1576  def color_by_uncertainty(self):
1577  self.method = self.color_by_uncertainty
1578  self.scheme = self.mpl.cm.jet
1579  ps = IMP.atom.get_leaves(self.hier)
1580  unc_dict = {}
1581  for p in ps:
1583  u = IMP.pmi.Uncertainty(p).get_uncertainty()
1584  unc_dict[p] = u
1585  self.first = self.get_log_scale(1.0)
1586  self.last = self.get_log_scale(100.0)
1587  for p in unc_dict:
1588  value = self.get_log_scale(unc_dict[p])
1589  if value >= self.last:
1590  value = self.last
1591  if value <= self.first:
1592  value = self.first
1593  c = self.get_color((value-self.first) / (self.last-self.first))
1594  IMP.display.Colored(p).set_color(c)
1595 
1596  def get_color_bar(self, filename):
1597  import matplotlib as mpl
1598  mpl.use('Agg')
1599  import matplotlib.pyplot as plt
1600  plt.clf()
1601  fig = plt.figure(figsize=(8, 3))
1602  ax1 = fig.add_axes([0.05, 0.80, 0.9, 0.15])
1603 
1604  cmap = self.scheme
1605  norm = mpl.colors.Normalize(vmin=0.0, vmax=1.0)
1606 
1607  if self.method == self.color_by_uncertainty:
1608  angticks = [1.0, 2.5, 5.0, 10.0, 25.0, 50.0, 100.0]
1609  vvalues = []
1610  marks = []
1611  for at in angticks:
1612  vvalue = (self.get_log_scale(at)-self.first) \
1613  / (self.last-self.first)
1614  if vvalue <= 1.0 and vvalue >= 0.0:
1615  vvalues.append(vvalue)
1616  marks.append(str(at))
1617  cb1 = mpl.colorbar.ColorbarBase(
1618  ax1, cmap=cmap, norm=norm, ticks=vvalues,
1619  orientation='horizontal')
1620  print(self.first, self.last, marks, vvalues)
1621  cb1.ax.set_xticklabels(marks)
1622  cb1.set_label('Angstorm')
1623  plt.savefig(filename, dpi=150, transparent=True)
1624  plt.show()
1625 
1626 
1627 def color2rgb(colorname):
1628  """Given a Chimera color name or hex color value, return RGB"""
1629  d = {'aquamarine': (0.4980392156862745, 1.0, 0.8313725490196079),
1630  'black': (0.0, 0.0, 0.0),
1631  'blue': (0.0, 0.0, 1.0),
1632  'brown': (0.6470588235, 0.16470588235294117, 0.16470588235294117),
1633  'chartreuse': (0.4980392156862745, 1.0, 0.0),
1634  'coral': (1.0, 0.4980392156862745, 0.3137254901960784),
1635  'cornflower blue': (0.39215686, 0.58431372549, 0.9294117647058824),
1636  'cyan': (0.0, 1.0, 1.0),
1637  'dark cyan': (0.0, 0.5450980392156862, 0.5450980392156862),
1638  'dark gray': (0.6627450980, 0.6627450980392157, 0.6627450980392157),
1639  'dark green': (0.0, 0.39215686274509803, 0.0),
1640  'dark khaki': (0.74117647, 0.7176470588235294, 0.4196078431372549),
1641  'dark magenta': (0.5450980392156862, 0.0, 0.5450980392156862),
1642  'dark olive green': (0.333333333, 0.419607843, 0.1843137254901961),
1643  'dark red': (0.5450980392156862, 0.0, 0.0),
1644  'dark slate blue': (0.28235294, 0.239215686, 0.5450980392156862),
1645  'dark slate gray': (0.1843137, 0.30980392, 0.30980392156862746),
1646  'deep pink': (1.0, 0.0784313725490196, 0.5764705882352941),
1647  'deep sky blue': (0.0, 0.7490196078431373, 1.0),
1648  'dim gray': (0.41176470, 0.4117647058823529, 0.4117647058823529),
1649  'dodger blue': (0.11764705882352941, 0.5647058823529412, 1.0),
1650  'firebrick': (0.6980392, 0.13333333333333333, 0.13333333333333333),
1651  'forest green': (0.13333333, 0.5450980392156862, 0.13333333333333333),
1652  'gold': (1.0, 0.8431372549019608, 0.0),
1653  'goldenrod': (0.85490196, 0.6470588235294118, 0.12549019607843137),
1654  'gray': (0.7450980392156863, 0.7450980392156863, 0.7450980392156863),
1655  'green': (0.0, 1.0, 0.0),
1656  'hot pink': (1.0, 0.4117647058823529, 0.7058823529411765),
1657  'khaki': (0.9411764705882353, 0.9019607843137255, 0.5490196078431373),
1658  'light blue': (0.67843137, 0.8470588235294118, 0.9019607843137255),
1659  'light gray': (0.82745098, 0.8274509803921568, 0.8274509803921568),
1660  'light green': (0.56470588, 0.9333333333333333, 0.5647058823529412),
1661  'light sea green': (0.125490, 0.6980392156862745, 0.6666666666666666),
1662  'lime green': (0.1960784, 0.803921568627451, 0.19607843137254902),
1663  'magenta': (1.0, 0.0, 1.0),
1664  'medium blue': (0.1960784, 0.19607843137254902, 0.803921568627451),
1665  'medium purple': (0.57647, 0.4392156862745098, 0.8588235294117647),
1666  'navy blue': (0.0, 0.0, 0.5019607843137255),
1667  'olive drab': (0.4196078, 0.5568627450980392, 0.13725490196078433),
1668  'orange red': (1.0, 0.27058823529411763, 0.0),
1669  'orange': (1.0, 0.4980392156862745, 0.0),
1670  'orchid': (0.85490196, 0.4392156862745098, 0.8392156862745098),
1671  'pink': (1.0, 0.7529411764705882, 0.796078431372549),
1672  'plum': (0.8666666666666667, 0.6274509803921569, 0.8666666666666667),
1673  'purple': (0.62745098, 0.12549019607843137, 0.9411764705882353),
1674  'red': (1.0, 0.0, 0.0),
1675  'rosy brown': (0.7372549, 0.5607843137254902, 0.5607843137254902),
1676  'salmon': (0.980392, 0.5019607843137255, 0.4470588235294118),
1677  'sandy brown': (0.956862745, 0.6431372549019608, 0.3764705882352941),
1678  'sea green': (0.18039, 0.5450980392156862, 0.3411764705882353),
1679  'sienna': (0.6274509, 0.3215686274509804, 0.17647058823529413),
1680  'sky blue': (0.52941176, 0.807843137254902, 0.9215686274509803),
1681  'slate gray': (0.439215686, 0.50196078, 0.5647058823529412),
1682  'spring green': (0.0, 1.0, 0.4980392156862745),
1683  'steel blue': (0.2745098, 0.50980392, 0.70588235),
1684  'tan': (0.8235294117647058, 0.7058823529411765, 0.5490196078431373),
1685  'turquoise': (0.25098039, 0.87843137, 0.81568627),
1686  'violet red': (0.81568627, 0.125490196, 0.56470588235),
1687  'white': (1.0, 1.0, 1.0),
1688  'yellow': (1.0, 1.0, 0.0)}
1689  if colorname.startswith('#'):
1690  return tuple(int(colorname[i:i+2], 16) / 255. for i in (1, 3, 5))
1691  else:
1692  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:660
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:1135
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:1106
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:743
def shuffle_configuration
Shuffle particles.
Definition: tools.py:1315
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:1173
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:1232
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:806
This class stores integers in ordered compact lists eg: [[1,2,3],[6,7,8]] the methods help splitting ...
Definition: tools.py:680
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:708
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:1128
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:1627
def input_adaptor
Adapt things for PMI (degrees of freedom, restraints, ...) Returns list of list of hierarchies...
Definition: tools.py:945
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:697
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:761
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:48
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:1279
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:919
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:1080
def get_rbs_and_beads
Returns unique objects in original order.
Definition: tools.py:1208
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: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:889
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