IMP  2.3.1
The Integrative Modeling Platform
em/analyze_convergence.py

Analyze the convergence of the IMP.em.FitRestraint. The script build a simple model and then displays the derivatives, em score and how well conjugate gradients converges under various displacements of the model.

1 ## \example em/analyze_convergence.py
2 # Analyze the convergence of the IMP.em.FitRestraint. The script build a
3 # simple model and then displays the derivatives, em score and how well
4 # conjugate gradients converges under various displacements of the model.
5 
6 import IMP.display
7 import IMP.em
8 
9 use_rigid_bodies = True
10 bd = 10
11 radius = 10
12 
13 m = IMP.kernel.Model()
14 p = IMP.kernel.Particle(m)
15 IMP.atom.Mass.setup_particle(p, 10000)
16 d = IMP.core.XYZR.setup_particle(p)
17 d.set_radius(radius)
18 
19 # Set up the particle as either a rigid body or a simple ball
20 if use_rigid_bodies:
21  prb = IMP.kernel.Particle(m)
22  prb.set_name("rigid body")
23  d.set_coordinates(IMP.algebra.Vector3D(0, 0, 0))
24  drb = IMP.core.RigidBody.setup_particle(
25  prb, IMP.algebra.ReferenceFrame3D())
26  drb.add_member(p)
27  print "initial frame", drb.get_reference_frame()
28  fp = prb
29  drb.set_coordinates_are_optimized(True)
30  refiner = IMP.core.TableRefiner()
31  refiner.add_particle(prb, [p])
32  to_move = drb
33  print [p.get_name() for p in refiner.get_refined(prb)]
34  fp = d
35 else:
36  fp = d
37  to_move = d
38  d.set_coordinates_are_optimized(True)
39  refiner = None
40 
41 
42 bb = IMP.algebra.BoundingBox3D(
43  IMP.algebra.Vector3D(-bd - radius, -bd - radius, -bd - radius),
44  IMP.algebra.Vector3D(bd + radius, bd + radius, bd + radius))
45 
46 dheader = IMP.em.create_density_header(bb, 1)
47 dheader.set_resolution(1)
48 dmap = IMP.em.SampledDensityMap(dheader)
49 dmap.set_particles([p])
50 
51 dmap.resample()
52 # computes statistic stuff about the map and insert it in the header
53 dmap.calcRMS()
54 IMP.em.write_map(dmap, "map.mrc", IMP.em.MRCReaderWriter())
55 rs = IMP.kernel.RestraintSet(m)
56 m.add_restraint(rs)
57 # rs.set_weight(.003)
58 
59 # if rigid bodies are used, we need to define a refiner as
60 # FitRestraint doesn't support just passing all the geometry
61 r = IMP.em.FitRestraint([fp], dmap)
62 rs.add_restraint(r)
63 g = IMP.core.XYZDerivativeGeometry(d)
64 g.set_name("deriv")
65 w = IMP.display.PymolWriter("derivatives.pym")
66 # kind of abusive
67 steps = 4
68 m.set_log_level(IMP.base.SILENT)
69 
70 opt = IMP.core.ConjugateGradients(m)
71 
72 
73 def try_point(i, j, k):
74  print "trying", i, j, k
75  vc = IMP.algebra.Vector3D(i, j, k)
76  to_move.set_coordinates(vc)
77  # display the score at this position
78  cg = IMP.display.SphereGeometry(IMP.algebra.Sphere3D(vc, 1))
79  cg.set_name("score")
80  v = m.evaluate(True)
81  cg.set_color(IMP.display.get_hot_color(v))
82  w.add_geometry(cg)
83  print "score and derivatives", v, to_move.get_derivatives()
84  w.add_geometry(g)
85 
86  opt.optimize(10)
87  print "after", d.get_coordinates()
88  mag = to_move.get_coordinates().get_magnitude()
89 
90  converge_color = IMP.display.get_gray_color(1.0 / (1.0 + mag))
91  # display the distance after optimization at this position
92  sg = IMP.display.SphereGeometry(IMP.algebra.Sphere3D(vc, 1))
93  sg.set_color(converge_color)
94  sg.set_name("converge")
95  w.add_geometry(sg)
96 
97 try_point(-bd, -bd, -bd)
98 
99 # For a more informative (but much slower) test, use the following instead:
100 # for i in range(-bd, bd+1, 2*bd/steps):
101 # for j in range(-bd, bd+1, 2*bd/steps):
102 # for k in range(-bd, bd+1, 2*bd/steps):
103 # try_point(i, j, k)