IMP  2.0.0
The Integrative Modeling Platform
restraints.py
1 
2 import IMP
3 import IMP.core as core
4 import IMP.domino as domino
5 import IMP.container as container
6 import IMP.em2d as em2d
7 import IMP.em2d.imp_general.representation as representation
8 import IMP.atom as atom
9 import IMP.base as base
10 import os
11 
12 import logging
13 
14 log = logging.getLogger("restraints")
15 
16 def get_connectivity_restraint(particles, distance=10.,
17  n_pairs=1, spring_constant=1):
18  """
19  Set a connectivity restraint for the leaves of a set of particles
20 
21  The intended use is that the each particle is a hierarchy. Each
22  hierarchy contains leaves that are atoms, or particles
23  that are a coarse representation of a molecule
24  """
25  score = core.HarmonicUpperBoundSphereDistancePairScore(distance,
26  spring_constant)
28  # score based on the one pair (the closest leaves)
29  pair_score = IMP.core.KClosePairsPairScore(score, refiner, n_pairs)
30  cr = IMP.core.ConnectivityRestraint(pair_score)
31  cr.set_particles(particles)
32  return cr
33 
34 
35 def get_em2d_restraint( assembly,
36  images_selection_file,
37  restraint_params,
38  mode="fast",
39  n_optimized=1):
40  """ Sets a restraint for comparing the model to a set of EM images
41  """
42  model = assembly.get_model()
43  # Setup the restraint
44  sc = em2d.EM2DScore()
45  r = em2d.Em2DRestraint()
46  r.setup(sc, restraint_params)
47  names = em2d.read_selection_file(images_selection_file)
48  names = [base.get_relative_path(images_selection_file, x) for x in names]
49  log.debug("names of the images %s", names)
50  srw = em2d.SpiderImageReaderWriter()
51  imgs = em2d.read_images(names, srw)
52  r.set_images(imgs)
53 
54  ps = atom.get_leaves(assembly)
55  lsc = container.ListSingletonContainer(ps)
56  r.set_particles(lsc)
57 
58  if (mode == "coarse"):
59  r.set_coarse_registration_mode(True)
60  elif (mode == "fast"):
61  r.set_fast_mode(n_optimized)
62  elif(mode == "complete"):
63  pass
64  else:
65  raise ValueError("Em2DRestraint mode not recognized")
66  return r