IMP logo
IMP Reference Guide  2.7.0
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 from __future__ import print_function
6 import IMP
7 import sys
8 
9 IMP.setup_from_argv(sys.argv, "Restraint example")
10 
11 # a restraint which checks if particles are sorted in
12 # increasing order on k.
13 
14 
15 class MyRestraint(IMP.Restraint):
16  # take the list of particles and the key to use
17 
18  def __init__(self, m, ps, k):
19  IMP.Restraint.__init__(self, m, "MyRestraint %1%")
20  self.ps = ps
21  self.k = k
22 
23  def unprotected_evaluate(self, da):
24  score = 0
25  for i in range(1, len(self.ps)):
26  p0 = self.ps[i - 1]
27  p1 = self.ps[i]
28  if p0.get_value(k) > p1.get_value(k):
29  diff = (p0.get_value(k) - p1.get_value(k))
30  score = score + diff
31  p0.add_to_derivative(k, -1, da)
32  p1.add_to_derivative(k, 1, da)
33  else:
34  if IMP.get_log_level() >= IMP.TERSE:
35  print(p0.get_name(), "and", p1.get_name(), " are ok")
36  return score
37 
38  def do_get_inputs(self):
39  return self.ps
40 
41 # some code to create and evaluate it
42 k = IMP.FloatKey("a key")
43 m = IMP.Model()
44 ps = []
45 for i in range(0, 10):
46  p = IMP.Particle(m)
47  p.add_attribute(k, i)
48  ps.append(p)
49 r = MyRestraint(m, ps, k)
50 # IMP.set_log_level(IMP.TERSE)
51 print(r.evaluate(True))