IMP  2.3.0
The Integrative Modeling Platform
em2d/optimize_em2d_with_montecarlo.py

Example of optimizing an EM2D restraint using Monte Carlo.

1 ## \example em2d/optimize_em2d_with_montecarlo.py
2 # Example of optimizing an EM2D restraint using Monte Carlo.
3 #
4 
5 import IMP
6 import IMP.base
7 import IMP.core
8 import IMP.atom
9 import IMP.em2d
10 import IMP.algebra
11 import IMP.container
12 import random
13 
14 
15 # An Optimizer score to get the values of the statistics after a given set
16 # of evaluations
17 class WriteStatisticsOptimizerScore(IMP.OptimizerState):
18 
19  def __init__(self, m, restraints):
20  IMP.OptimizerState.__init__(self, m, "WriteStats")
21  self.count = 0
22  self.restraints = restraints
23 
24  def update(self):
25  if (self.count != 10):
26  self.count += 1
27  return
28  else:
29  self.count = 0
30  o = self.get_optimizer()
31  m = o.get_model()
32  for r in self.restraints:
33  print r.get_name(), r.get_last_score()
34  # for i in range(0,m.get_number_of_restraints()):
35  # r=m.get_restraint(i)
36  # print "restraint",r.get_name(),"value",r.evaluate(False)
37 
38 
39 # Get model from PDB file
40 IMP.base.set_log_level(IMP.base.TERSE)
41 m = IMP.kernel.Model()
42 prot = IMP.atom.read_pdb(
43  IMP.em2d.get_example_path("1z5s.pdb"), m, IMP.atom.ATOMPDBSelector())
44 IMP.atom.add_radii(prot)
45 
46 # get the chains
47 chains = IMP.atom.get_by_type(prot, IMP.atom.CHAIN_TYPE)
48 print "there are", len(chains), "chains in 1z5s.pdb"
49 
50 # set the chains as rigid bodies
51 native_chain_centers = []
52 rigid_bodies = []
53 for c in chains:
54  atoms = IMP.core.get_leaves(c)
55  rbd = IMP.core.RigidBody.setup_particle(c, atoms)
56  rbd.set_coordinates_are_optimized(True)
57  rigid_bodies.append(rbd)
58  print "chain has", rbd.get_number_of_members(), \
59  "atoms", "coordinates: ", rbd.get_coordinates()
60  native_chain_centers.append(rbd.get_coordinates())
61 
62 bb = IMP.algebra.BoundingBox3D(IMP.algebra.Vector3D(-25, -40, -60),
63  IMP.algebra.Vector3D(25, 40, 60))
64 # rotate and translate the chains
65 for rbd in rigid_bodies:
66  # rotation
67  rotation = IMP.algebra.get_random_rotation_3d()
68  transformation1 = IMP.algebra.get_rotation_about_point(
69  rbd.get_coordinates(), rotation)
70  # translation
71  transformation2 = IMP.algebra.Transformation3D(
72  IMP.algebra.get_random_vector_in(bb))
73  # Apply
74  final_transformation = IMP.algebra.compose(
75  transformation1, transformation2)
76  IMP.core.transform(rbd, final_transformation)
77 print "Writing transformed assembly"
78 IMP.atom.write_pdb(prot, "1z5s-transformed.pdb")
79 
80 # set distance restraints measusring some distances between rigid bodies
81 # for the solution.
82 d01 = IMP.algebra.get_distance(
83  native_chain_centers[0], native_chain_centers[1])
84 r01 = IMP.core.DistanceRestraint(
85  IMP.core.Harmonic(d01, 1), chains[0], chains[1])
86 r01.set_name("distance 0-1")
87 d12 = IMP.algebra.get_distance(
88  native_chain_centers[1], native_chain_centers[2])
89 r12 = IMP.core.DistanceRestraint(
90  IMP.core.Harmonic(d12, 1), chains[1], chains[2])
91 r12.set_name("distance 1-2")
92 d23 = IMP.algebra.get_distance(
93  native_chain_centers[2], native_chain_centers[3])
94 r23 = IMP.core.DistanceRestraint(
95  IMP.core.Harmonic(d23, 1), chains[2], chains[3])
96 r23.set_name("distance 2-3")
97 d30 = IMP.algebra.get_distance(
98  native_chain_centers[3], native_chain_centers[0])
99 r30 = IMP.core.DistanceRestraint(
100  IMP.core.Harmonic(d30, 1), chains[3], chains[0])
101 r30.set_name("distance 3-0")
102 print "Distances in the solution: d01 =", \
103  d01, "d12 =", d12, "d23 =", d23, "d30 =", d30
104 
105 # set distance restraints
106 print "adding distance restraints "
107 for r in [r01, r12, r23, r30]:
108  m.add_restraint(r)
109 print "model has ", m.get_number_of_restraints(), "restraints"
110 
111 # set em2D restraint
112 srw = IMP.em2d.SpiderImageReaderWriter()
113 selection_file = IMP.em2d.get_example_path("all-1z5s-projections.sel")
114 images_to_read_names = [IMP.base.get_relative_path(selection_file, x) for x in
115  IMP.em2d.read_selection_file(selection_file)]
116 em_images = IMP.em2d.read_images(images_to_read_names, srw)
117 print len(em_images), "images read"
118 
119 em2d_restraint = IMP.em2d.Em2DRestraint(m)
120 apix = 1.5 # sampling rate of the available EM images
121 # resolution at which you want to generate the projections of the model
122 # In principle you want "perfect" projections, so use the highest resolution
123 resolution = 1
124 # Number of projections to use for the initial registration
125 # (coarse registration) to estimate the registration parameters
126 n_projections = 20
127 params = IMP.em2d.Em2DRestraintParameters(apix, resolution, n_projections)
128 
129 # This method (recommended) uses preprocessing of the images and projections
130 # to speed-up the registration
131 params.coarse_registration_method = IMP.em2d.ALIGN2D_PREPROCESSING
132 # use true if you want to save the projections from the model that best
133 # match the Em images
134 params.save_match_images = False
135 
136 score_function = IMP.em2d.EM2DScore()
137 em2d_restraint = IMP.em2d.Em2DRestraint(m)
138 em2d_restraint.setup(score_function, params)
139 em2d_restraint.set_images(em_images)
140 em2d_restraint.set_name("em2d restraint")
141 container = IMP.container.ListSingletonContainer(IMP.core.get_leaves(prot))
142 em2d_restraint.set_particles(container)
143 em2d_restraints_set = IMP.kernel.RestraintSet(m)
144 
145 # The next two lines are commented, because the optimization of the example
146 # is expensive. To run the full example, uncomment them (It can take a few
147 # hours).
148 # em2d_restraints_set.add_restraint(em2d_restraint)
149 # em2d_restraints_set.set_weight(1000) # weight for the em2D restraint
150 
151 print "adding em2d restraint "
152 m.add_restraint(em2d_restraints_set)
153 # Add all restraints to a model
154 print "model has ", m.get_number_of_restraints(), "restraints"
155 
156 
157 # MONTECARLO OPTIMIZATION
158 s = IMP.core.MonteCarlo(m)
159 # Add movers for the rigid bodies
160 movers = []
161 for rbd in rigid_bodies:
162  movers.append(IMP.core.RigidBodyMover(rbd, 5, 2))
163 s.add_movers(movers)
164 print "MonteCarlo sampler has", s.get_number_of_movers(), "movers"
165 # Add an optimizer state to save intermediate configurations of the hierarchy
166 o_state = IMP.atom.WritePDBOptimizerState(chains, "intermediate-step-%1%.pdb")
167 o_state.set_period(10)
168 s.add_optimizer_state(o_state)
169 
170 ostate2 = WriteStatisticsOptimizerScore(m, m.get_restraints())
171 s.add_optimizer_state(ostate2)
172 
173 # Perform optimization
174 temperatures = [200, 100, 60, 40, 20, 5]
175 # 200 steps recommended for accurate optimization; a smaller number is used
176 # here for demonstration purposes
177 optimization_steps = 10
178 for T in temperatures:
179  s.set_kt(T)
180  s.optimize(optimization_steps)
181 IMP.atom.write_pdb(prot, "solution.pdb")
182 
183 
184 # Check that the optimization achieves distances close to the solution
185 print "*** End optimization ***"
186 new_centers = []
187 for rbd in rigid_bodies:
188  print "chain has", rbd.get_number_of_members(), \
189  "atoms", "coordinates: ", rbd.get_coordinates()
190  new_centers.append(rbd.get_coordinates())
191 
192 d01 = IMP.algebra.get_distance(new_centers[0], new_centers[1])
193 d12 = IMP.algebra.get_distance(new_centers[1], new_centers[2])
194 d23 = IMP.algebra.get_distance(new_centers[2], new_centers[3])
195 d30 = IMP.algebra.get_distance(new_centers[3], new_centers[0])
196 print "Distances at the end of the optimization: d01 =", \
197  d01, "d12 =", d12, "d23 =", d23, "d30 =", d30