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