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