IMP  2.0.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 class MyRestraint(IMP.Restraint):
10  # take the list of particles and the key to use
11  def __init__(self, m, ps, k):
12  IMP.Restraint.__init__(self, m)
13  self.ps=ps
14  self.k=k
15  def unprotected_evaluate(self, da):
16  score=0
17  for i in range(1, len(self.ps)):
18  p0= self.ps[i-1]
19  p1= self.ps[i]
20  if p0.get_value(k) > p1.get_value(k):
21  diff=(p0.get_value(k)-p1.get_value(k))
22  score= score+diff
23  p0.add_to_derivative(k, -1, da)
24  p1.add_to_derivative(k, 1, da)
25  else:
26  if IMP.get_log_level() >= IMP.base.TERSE:
27  print p0.get_name(), "and", p1.get_name(), " are ok"
28  return score
29  def get_input_particles(self):
30  return self.ps
31  def get_input_containers(self):
32  return []
33 
34 # some code to create and evaluate it
35 k= IMP.FloatKey("a key")
36 m= IMP.Model()
37 ps=[]
38 for i in range(0,10):
39  p = IMP.Particle(m)
40  p.add_attribute(k, i)
41  ps.append(p)
42 r= MyRestraint(m, ps, k)
43 #IMP.base.set_log_level(IMP.base.TERSE)
44 print r.evaluate(True)