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