IMP  2.1.1
The Integrative Modeling Platform
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 (incremental) and conjugate gradients in conjunction in a non-trivial optimization.

1 ## \example core/optimize_balls.py
2 # This example optimizes a set of a balls to form 100 chains packed into a
3 # box. It illustrates using Monte Carlo (incremental) and conjugate
4 # gradients in conjunction in a non-trivial optimization.
5 
6 import IMP.core
7 import IMP.display
8 import IMP.container
9 import IMP.rmf
10 import IMP.base
11 import sys
12 
13 IMP.base.setup_from_argv(sys.argv, "Optimize balls example")
14 
15 bb = IMP.algebra.BoundingBox3D(IMP.algebra.Vector3D(0, 0, 0),
16  IMP.algebra.Vector3D(10, 10, 10))
17 # in fast do 10,10,10, for the purposes of testing we reduce it
18 ni = 2
19 nj = 2
20 np = 2
21 radius = .45
22 k = 100
23 
24 # using a HarmonicDistancePairScore for fixed length links is more
25 # efficient than using a HarmonicSphereDistnacePairScore and works
26 # better with the optimizer
27 lps = IMP.core.HarmonicDistancePairScore(1.5 * radius, k)
28 sps = IMP.core.SoftSpherePairScore(k)
29 
30 m = IMP.kernel.Model()
31 IMP.base.set_log_level(IMP.base.SILENT)
32 aps = []
33 filters = []
34 movers = []
35 restraints = []
36 for i in range(0, ni):
37  for j in range(0, nj):
38  base = IMP.algebra.Vector3D(i, j, 0)
39  chain = []
40  for k in range(0, np):
41  p = IMP.kernel.Particle(m)
42  p.set_name("P" + str(i) + " " + str(j) + " " + str(k))
43  s = IMP.algebra.Sphere3D(
44  IMP.algebra.get_random_vector_in(bb), radius)
45  d = IMP.core.XYZR.setup_particle(p, s)
46  d.set_coordinates_are_optimized(True)
47  movers.append(IMP.core.BallMover([p], radius * 2))
48  movers[-1].set_was_used(True)
49  IMP.display.Colored.setup_particle(
50  p, IMP.display.get_display_color(i * nj + j))
51  if k == 0:
52  d.set_coordinates(base)
53  else:
54  d.set_coordinates_are_optimized(True)
55  chain.append(p)
56  aps.append(p)
57  # set up a chain of bonds
58  cpc = IMP.container.ExclusiveConsecutivePairContainer(chain)
59  r = IMP.container.PairsRestraint(lps, cpc)
60  restraints.append(r)
61 
62 # don't apply excluded volume to consecutive particles
63 filters.append(IMP.container.ExclusiveConsecutivePairFilter())
64 ibss = IMP.core.BoundingBox3DSingletonScore(
65  IMP.core.HarmonicUpperBound(0, k), bb)
66 bbr = IMP.container.SingletonsRestraint(ibss, aps)
67 restraints.append(bbr)
68 
69 cg = IMP.core.ConjugateGradients(m)
70 mc = IMP.core.MonteCarlo(m)
71 mc.set_name("MC")
72 sm = IMP.core.SerialMover(movers)
73 mc.add_mover(sm)
74 # we are special casing the nbl term
75 isf = IMP.core.IncrementalScoringFunction(aps, restraints)
76 isf.set_name("I")
77 # use special incremental support for the non-bonded part
78 # apply the pair score sps to all touching ball pairs from the list of particles
79 # aps, using the filters to remove undersired pairs
80 # this is equivalent to the nbl construction above but optimized for
81 # incremental
82 isf.add_close_pair_score(sps, 0, aps, filters)
83 
84 # create a scoring function for conjugate gradients that includes the
85 # ExcludedVolumeRestraint
86 nbl = IMP.core.ExcludedVolumeRestraint(aps, k, 1)
87 nbl.set_pair_filters(filters)
88 sf = IMP.core.RestraintsScoringFunction(restraints + [nbl], "RSF")
89 
90 if True:
91  mc.set_incremental_scoring_function(isf)
92 else:
93  # we could, instead do non-incremental scoring
94  mc.set_scoring_function(sf)
95 
96 # first relax the bonds a bit
97 rs = []
98 for p in aps:
99  rs.append(IMP.ScopedSetFloatAttribute(p, IMP.core.XYZR.get_radius_key(),
100  0))
101 cg.set_scoring_function(sf)
102 cg.optimize(1000)
103 for r in restraints:
104  print r.get_name(), r.evaluate(False)
105 
106 # shrink each of the particles, relax the configuration, repeat
107 for i in range(1, 11):
108  rs = []
109  factor = .1 * i
110  for p in aps:
111  rs.append(
112  IMP.ScopedSetFloatAttribute(p, IMP.core.XYZR.get_radius_key(),
113  IMP.core.XYZR(p).get_radius() * factor))
114  # move each particle 100 times
115  print factor
116  for j in range(0, 5):
117  print "stage", j
118  mc.set_kt(100.0 / (3 * j + 1))
119  print "mc", mc.optimize(ni * nj * np * (j + 1) * 100), cg.optimize(10)
120  del rs
121  for r in restraints:
122  print r.get_name(), r.evaluate(False)
123 
124 w = IMP.display.PymolWriter("final.pym")
125 for p in aps:
126  g = IMP.core.XYZRGeometry(p)
127  w.add_geometry(g)
128 g = IMP.display.BoundingBoxGeometry(bb)
129 g.set_name("bb")
130 w.add_geometry(g)