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