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