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