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