/
optimize_balls.py
135 lines (122 loc) · 4.02 KB
/
optimize_balls.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
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
## \example core/optimize_balls.py
# This example optimizes a set of a balls to form 100 chains packed into a
# box. It illustrates using Monte Carlo and conjugate
# gradients in conjunction in a non-trivial optimization.
import IMP.core
import IMP.display
import IMP.container
import sys
IMP.setup_from_argv(sys.argv, "Optimize balls example")
bb = IMP.algebra.BoundingBox3D(IMP.algebra.Vector3D(0, 0, 0),
IMP.algebra.Vector3D(10, 10, 10))
if IMP.get_is_quick_test():
ni = 2
nj = 2
np = 2
radius = .45
k = 100
ncg = 10
nmc = 1
ninner = 1
nouter = 1
else:
ni = 10
nj = 10
np = 10
radius = .45
k = 100
ncg = 1000
nmc = ni * nj * np * 100
ninner = 5
nouter = 11
print(IMP.get_is_quick_test(), ni, nj, np, ninner, nouter)
# using a HarmonicDistancePairScore for fixed length links is more
# efficient than using a HarmonicSphereDistnacePairScore and works
# better with the optimizer
lps = IMP.core.HarmonicDistancePairScore(1.5 * radius, k)
sps = IMP.core.SoftSpherePairScore(k)
m = IMP.Model()
# IMP.set_log_level(IMP.SILENT)
aps = []
filters = []
movers = []
restraints = []
for i in range(0, ni):
for j in range(0, nj):
base = IMP.algebra.Vector3D(i, j, 0)
chain = []
for k in range(0, np):
p = IMP.Particle(m)
p.set_name("P" + str(i) + " " + str(j) + " " + str(k))
s = IMP.algebra.Sphere3D(
IMP.algebra.get_random_vector_in(bb), radius)
d = IMP.core.XYZR.setup_particle(p, s)
d.set_coordinates_are_optimized(True)
movers.append(IMP.core.BallMover(m, p, radius * 2))
movers[-1].set_was_used(True)
IMP.display.Colored.setup_particle(
p, IMP.display.get_display_color(i * nj + j))
if k == 0:
d.set_coordinates(base)
else:
d.set_coordinates_are_optimized(True)
chain.append(p)
aps.append(p)
# set up a chain of bonds
cpc = IMP.container.ExclusiveConsecutivePairContainer(m, chain)
r = IMP.container.PairsRestraint(lps, cpc)
restraints.append(r)
# don't apply excluded volume to consecutive particles
filters.append(IMP.container.ExclusiveConsecutivePairFilter())
ibss = IMP.core.BoundingBox3DSingletonScore(
IMP.core.HarmonicUpperBound(0, k), bb)
bbr = IMP.container.SingletonsRestraint(ibss, aps)
restraints.append(bbr)
cg = IMP.core.ConjugateGradients(m)
mc = IMP.core.MonteCarlo(m)
mc.set_name("MC")
sm = IMP.core.SerialMover(movers)
mc.add_mover(sm)
# create a scoring function for conjugate gradients that includes the
# ExcludedVolumeRestraint
nbl = IMP.core.ExcludedVolumeRestraint(aps, k, 1)
nbl.set_pair_filters(filters)
sf = IMP.core.RestraintsScoringFunction(restraints + [nbl], "RSF")
mc.set_scoring_function(sf)
# Speed up the optimization by only rescoring terms involving particles
# that moved at each MC step
mc.set_score_moved(True)
# first relax the bonds a bit
rs = []
for p in aps:
rs.append(IMP.ScopedSetFloatAttribute(p, IMP.core.XYZR.get_radius_key(),
0))
cg.set_scoring_function(sf)
cg.optimize(ncg)
for r in restraints:
print(r.get_name(), r.evaluate(False))
# shrink each of the particles, relax the configuration, repeat
for i in range(1, nouter):
rs = []
factor = .1 * i
for p in aps:
rs.append(
IMP.ScopedSetFloatAttribute(
p, IMP.core.XYZR.get_radius_key(),
IMP.core.XYZR(p).get_radius() * factor))
# move each particle nmc times
print(factor)
for j in range(0, ninner):
print("stage", j)
mc.set_kt(100.0 / (3 * j + 1))
print("mc", mc.optimize((j + 1) * nmc), cg.optimize(nmc))
del rs
for r in restraints:
print(r.get_name(), r.evaluate(False))
w = IMP.display.PymolWriter("final.pym")
for p in aps:
g = IMP.core.XYZRGeometry(p)
w.add_geometry(g)
g = IMP.display.BoundingBoxGeometry(bb)
g.set_name("bb")
w.add_geometry(g)