IMP  2.3.1
The Integrative Modeling Platform
kernel/write_a_restraint.py

While we do not recommend doing serious work using restraints written in Python, it is often useful when prototyping or testing code. Copy this example and modify as needed.

1 ## \example kernel/write_a_restraint.py
2 # While we do not recommend doing serious work using restraints written in Python, it is often useful when prototyping or testing code. Copy this example and modify as needed.
3 #
4 
5 import IMP
6 
7 # a restraint which checks if particles are sorted in
8 # increasing order on k.
9 
10 
11 class MyRestraint(IMP.kernel.Restraint):
12  # take the list of particles and the key to use
13 
14  def __init__(self, m, ps, k):
15  IMP.kernel.Restraint.__init__(self, m, "MyRestraint %1%")
16  self.ps = ps
17  self.k = k
18 
19  def unprotected_evaluate(self, da):
20  score = 0
21  for i in range(1, len(self.ps)):
22  p0 = self.ps[i - 1]
23  p1 = self.ps[i]
24  if p0.get_value(k) > p1.get_value(k):
25  diff = (p0.get_value(k) - p1.get_value(k))
26  score = score + diff
27  p0.add_to_derivative(k, -1, da)
28  p1.add_to_derivative(k, 1, da)
29  else:
30  if IMP.get_log_level() >= IMP.base.TERSE:
31  print p0.get_name(), "and", p1.get_name(), " are ok"
32  return score
33 
34  def do_get_inputs(self):
35  return self.ps
36 
37 # some code to create and evaluate it
38 k = IMP.FloatKey("a key")
39 m = IMP.kernel.Model()
40 ps = []
41 for i in range(0, 10):
42  p = IMP.kernel.Particle(m)
43  p.add_attribute(k, i)
44  ps.append(p)
45 r = MyRestraint(m, ps, k)
46 # IMP.base.set_log_level(IMP.base.TERSE)
47 print r.evaluate(True)