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