1
"""
2
Tests the restraints utilities.
3
"""
4

5 26
import logging
6 26
import os
7

8 26
import numpy as np
9 26
import parmed as pmd
10 26
import pytest
11

12 26
from paprika.restraints.restraints import DAT_restraint, create_window_list
13 26
from paprika.restraints.utils import (
14
    extract_guest_restraints,
15
    get_bias_potential_type,
16
    get_restraint_values,
17
)
18

19 26
logger = logging.getLogger(__name__)
20

21

22 26
def test_DAT_restraint():
23
    # Method 1
24 26
    logger.info("### Testing restraint 1, Method 1")
25 26
    rest1 = DAT_restraint()
26 26
    rest1.amber_index = True
27 26
    rest1.continuous_apr = False
28 26
    rest1.auto_apr = False
29 26
    rest1.topology = os.path.join(
30
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
31
    )
32 26
    rest1.mask1 = ":CB6@O,O2,O4,O6,O8,O10"
33 26
    rest1.mask2 = ":BUT@C3"
34 26
    rest1.attach["target"] = 3.0
35 26
    rest1.attach["num_windows"] = 4
36 26
    rest1.attach["fc_initial"] = 0.0
37 26
    rest1.attach["fc_final"] = 3.0
38 26
    rest1.pull["fc"] = rest1.attach["fc_final"]
39 26
    rest1.pull["num_windows"] = 4
40 26
    rest1.pull["target_initial"] = rest1.attach["target"]
41 26
    rest1.pull["target_final"] = 6.0
42 26
    rest1.release["target"] = rest1.pull["target_final"]
43 26
    rest1.release["num_windows"] = rest1.attach["num_windows"]
44 26
    rest1.release["fc_initial"] = rest1.attach["fc_initial"]
45 26
    rest1.release["fc_final"] = rest1.attach["fc_final"]
46 26
    rest1.initialize()
47 26
    assert rest1.index1 == [13, 31, 49, 67, 85, 103]
48 26
    assert rest1.index2 == [119]
49 26
    assert rest1.index3 is None
50 26
    assert rest1.index4 is None
51 26
    assert np.allclose(
52
        rest1.phase["attach"]["force_constants"], np.array([0.0, 1.0, 2.0, 3.0])
53
    )
54 26
    assert np.allclose(rest1.phase["attach"]["targets"], np.array([3.0, 3.0, 3.0, 3.0]))
55 26
    assert np.allclose(
56
        rest1.phase["pull"]["force_constants"], np.array([3.0, 3.0, 3.0, 3.0])
57
    )
58 26
    assert np.allclose(rest1.phase["pull"]["targets"], np.array([3.0, 4.0, 5.0, 6.0]))
59 26
    assert np.allclose(
60
        rest1.phase["release"]["force_constants"], np.array([0.0, 1.0, 2.0, 3.0])
61
    )
62 26
    assert np.allclose(
63
        rest1.phase["release"]["targets"], np.array([6.0, 6.0, 6.0, 6.0])
64
    )
65 26
    window_list = create_window_list([rest1])
66 26
    assert window_list == [
67
        "a000",
68
        "a001",
69
        "a002",
70
        "a003",
71
        "p000",
72
        "p001",
73
        "p002",
74
        "p003",
75
        "r000",
76
        "r001",
77
        "r002",
78
        "r003",
79
    ]
80

81
    # Method 1a
82 26
    logger.info("### Testing restraint 2, Method 1a")
83 26
    rest2 = DAT_restraint()
84 26
    rest2.amber_index = True
85 26
    rest2.continuous_apr = False
86 26
    rest2.auto_apr = False
87 26
    rest2.topology = os.path.join(
88
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
89
    )
90 26
    rest2.mask1 = ":CB6@O,O2,O4,O6,O8,O10"
91 26
    rest2.mask2 = ":BUT@C3"
92 26
    rest2.mask3 = ":BUT@C"
93 26
    rest2.attach["target"] = 180.0
94 26
    rest2.attach["num_windows"] = 4
95 26
    rest2.attach["fc_final"] = 75.0
96 26
    rest2.pull["fc"] = rest2.attach["fc_final"]
97 26
    rest2.pull["num_windows"] = 4
98 26
    rest2.pull["target_final"] = 180.0
99 26
    rest2.release["target"] = rest2.pull["target_final"]
100 26
    rest2.release["num_windows"] = rest2.attach["num_windows"]
101 26
    rest2.release["fc_final"] = rest2.attach["fc_final"]
102 26
    rest2.initialize()
103 26
    assert rest2.index1 == [13, 31, 49, 67, 85, 103]
104 26
    assert rest2.index2 == [119]
105 26
    assert rest2.index3 == [109]
106 26
    assert rest2.index4 is None
107 26
    assert np.allclose(
108
        rest2.phase["attach"]["force_constants"], np.array([0.0, 25.0, 50.0, 75.0])
109
    )
110 26
    assert np.allclose(
111
        rest2.phase["attach"]["targets"], np.array([180.0, 180.0, 180.0, 180.0])
112
    )
113 26
    assert np.allclose(
114
        rest2.phase["pull"]["force_constants"], np.array([75.0, 75.0, 75.0, 75.0])
115
    )
116 26
    assert np.allclose(
117
        rest2.phase["pull"]["targets"], np.array([0.0, 60.0, 120.0, 180.0])
118
    )
119 26
    assert np.allclose(
120
        rest2.phase["release"]["force_constants"], np.array([0.0, 25.0, 50.0, 75.0])
121
    )
122 26
    assert np.allclose(
123
        rest2.phase["release"]["targets"], np.array([180.0, 180.0, 180.0, 180.0])
124
    )
125 26
    window_list = create_window_list([rest2])
126 26
    assert window_list == [
127
        "a000",
128
        "a001",
129
        "a002",
130
        "a003",
131
        "p000",
132
        "p001",
133
        "p002",
134
        "p003",
135
        "r000",
136
        "r001",
137
        "r002",
138
        "r003",
139
    ]
140

141
    # Method 2 (Note auto_apr = True)
142 26
    logger.info("### Testing restraint 3, Method 2")
143 26
    rest3 = DAT_restraint()
144 26
    rest3.amber_index = True
145 26
    rest3.continuous_apr = False
146 26
    rest3.auto_apr = True
147 26
    rest3.topology = os.path.join(
148
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
149
    )
150 26
    rest3.mask1 = ":CB6@O2"
151 26
    rest3.mask2 = ":CB6@O"
152 26
    rest3.mask3 = ":BUT@C3"
153 26
    rest3.mask4 = ":BUT@C"
154 26
    rest3.attach["target"] = 90.0
155 26
    rest3.attach["fc_increment"] = 25.0
156 26
    rest3.attach["fc_initial"] = 0.0
157 26
    rest3.attach["fc_final"] = 75.0
158 26
    rest3.pull["target_increment"] = 1.0
159 26
    rest3.pull["target_final"] = 93.0
160 26
    rest3.release["fc_final"] = 75.0
161 26
    rest3.initialize()
162 26
    assert rest3.index1 == [31]
163 26
    assert rest3.index2 == [13]
164 26
    assert rest3.index3 == [119]
165 26
    assert rest3.index4 == [109]
166 26
    assert np.allclose(
167
        rest3.phase["attach"]["force_constants"], np.array([0.0, 25.0, 50.0, 75.0])
168
    )
169 26
    assert np.allclose(
170
        rest3.phase["attach"]["targets"], np.array([90.0, 90.0, 90.0, 90.0])
171
    )
172 26
    assert np.allclose(
173
        rest3.phase["pull"]["force_constants"], np.array([75.0, 75.0, 75.0, 75.0])
174
    )
175 26
    assert np.allclose(
176
        rest3.phase["pull"]["targets"], np.array([90.0, 91.0, 92.0, 93.0])
177
    )
178 26
    assert np.allclose(
179
        rest3.phase["release"]["force_constants"], np.array([0.0, 25.0, 50.0, 75.0])
180
    )
181 26
    assert np.allclose(
182
        rest3.phase["release"]["targets"], np.array([93.0, 93.0, 93.0, 93.0])
183
    )
184 26
    window_list = create_window_list([rest3])
185 26
    assert window_list == [
186
        "a000",
187
        "a001",
188
        "a002",
189
        "a003",
190
        "p000",
191
        "p001",
192
        "p002",
193
        "p003",
194
        "r000",
195
        "r001",
196
        "r002",
197
        "r003",
198
    ]
199

200
    # Method 2a
201 26
    logger.info("### Testing restraint 4, Method 2a")
202 26
    rest4 = DAT_restraint()
203 26
    rest4.amber_index = True
204 26
    rest4.continuous_apr = False
205 26
    rest4.auto_apr = False
206 26
    rest4.topology = os.path.join(
207
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
208
    )
209 26
    rest4.mask1 = ":CB6@O2"
210 26
    rest4.mask2 = ":CB6@O"
211 26
    rest4.mask3 = ":BUT@C3"
212 26
    rest4.mask4 = ":BUT@C"
213 26
    rest4.attach["target"] = 0.0
214 26
    rest4.attach["fc_increment"] = 25.0
215 26
    rest4.attach["fc_final"] = 75.0
216 26
    rest4.pull["fc"] = 75.0
217 26
    rest4.pull["target_increment"] = 1.0
218 26
    rest4.pull["target_final"] = 3.0
219 26
    rest4.release["target"] = 3.0
220 26
    rest4.release["fc_increment"] = 25.0
221 26
    rest4.release["fc_final"] = 75.0
222 26
    rest4.initialize()
223 26
    assert rest4.index1 == [31]
224 26
    assert rest4.index2 == [13]
225 26
    assert rest4.index3 == [119]
226 26
    assert rest4.index4 == [109]
227 26
    assert np.allclose(
228
        rest4.phase["attach"]["force_constants"], np.array([0.0, 25.0, 50.0, 75.0])
229
    )
230 26
    assert np.allclose(rest4.phase["attach"]["targets"], np.array([0.0, 0.0, 0.0, 0.0]))
231 26
    assert np.allclose(
232
        rest4.phase["pull"]["force_constants"], np.array([75.0, 75.0, 75.0, 75.0])
233
    )
234 26
    assert np.allclose(rest4.phase["pull"]["targets"], np.array([0.0, 1.0, 2.0, 3.0]))
235 26
    assert np.allclose(
236
        rest4.phase["release"]["force_constants"], np.array([0.0, 25.0, 50.0, 75.0])
237
    )
238 26
    assert np.allclose(
239
        rest4.phase["release"]["targets"], np.array([3.0, 3.0, 3.0, 3.0])
240
    )
241 26
    window_list = create_window_list([rest4])
242 26
    assert window_list == [
243
        "a000",
244
        "a001",
245
        "a002",
246
        "a003",
247
        "p000",
248
        "p001",
249
        "p002",
250
        "p003",
251
        "r000",
252
        "r001",
253
        "r002",
254
        "r003",
255
    ]
256

257
    # Method 3
258 26
    logger.info("### Testing restraint 5, Method 3")
259 26
    rest5 = DAT_restraint()
260 26
    rest5.amber_index = True
261 26
    rest5.continuous_apr = False
262 26
    rest5.auto_apr = False
263 26
    rest5.topology = os.path.join(
264
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
265
    )
266 26
    rest5.mask1 = ":CB6@O,O2,O4,O6,O8,O10"
267 26
    rest5.mask2 = ":BUT@C*"
268 26
    rest5.attach["target"] = 0.0
269 26
    rest5.attach["fraction_list"] = [0.0, 0.2, 0.5, 1.0]
270 26
    rest5.attach["fc_final"] = 5.0
271 26
    rest5.pull["fc"] = rest5.attach["fc_final"]
272 26
    rest5.pull["fraction_list"] = [0.0, 0.5, 1.0]
273 26
    rest5.pull["target_final"] = 1.0
274 26
    rest5.release["target"] = rest5.pull["target_final"]
275 26
    rest5.release["fraction_list"] = [0.0, 0.3, 0.6, 1.0]
276 26
    rest5.release["fc_final"] = rest5.attach["fc_final"]
277 26
    rest5.initialize()
278 26
    assert rest5.index1 == [13, 31, 49, 67, 85, 103]
279 26
    assert rest5.index2 == [109, 113, 115, 119]
280 26
    assert rest5.index3 is None
281 26
    assert rest5.index4 is None
282 26
    assert np.allclose(
283
        rest5.phase["attach"]["force_constants"], np.array([0.0, 1.0, 2.5, 5.0])
284
    )
285 26
    assert np.allclose(rest5.phase["attach"]["targets"], np.array([0.0, 0.0, 0.0, 0.0]))
286 26
    assert np.allclose(
287
        rest5.phase["pull"]["force_constants"], np.array([5.0, 5.0, 5.0])
288
    )
289 26
    assert np.allclose(rest5.phase["pull"]["targets"], np.array([0.0, 0.5, 1.0]))
290 26
    assert np.allclose(
291
        rest5.phase["release"]["force_constants"], np.array([0.0, 1.5, 3.0, 5.0])
292
    )
293 26
    assert np.allclose(
294
        rest5.phase["release"]["targets"], np.array([1.0, 1.0, 1.0, 1.0])
295
    )
296 26
    window_list = create_window_list([rest5])
297 26
    assert window_list == [
298
        "a000",
299
        "a001",
300
        "a002",
301
        "a003",
302
        "p000",
303
        "p001",
304
        "p002",
305
        "r000",
306
        "r001",
307
        "r002",
308
        "r003",
309
    ]
310

311
    # Method 4
312 26
    logger.info("### Testing restraint 6, Method 4")
313 26
    rest6 = DAT_restraint()
314 26
    rest6.amber_index = True
315 26
    rest6.continuous_apr = False
316 26
    rest6.auto_apr = False
317 26
    rest6.topology = os.path.join(
318
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
319
    )
320 26
    rest6.mask1 = ":CB6@O,O2,O4,O6,O8,O10"
321 26
    rest6.mask2 = ":BUT@C*"
322 26
    rest6.attach["target"] = 0.0
323 26
    rest6.attach["fraction_increment"] = 0.25
324 26
    rest6.attach["fc_final"] = 5.0
325 26
    rest6.pull["fc"] = rest6.attach["fc_final"]
326 26
    rest6.pull["fraction_increment"] = 0.5
327 26
    rest6.pull["target_final"] = 1.0
328 26
    rest6.release["target"] = rest6.pull["target_final"]
329 26
    rest6.release["fraction_increment"] = 0.33
330 26
    rest6.release["fc_final"] = rest6.attach["fc_final"]
331 26
    rest6.initialize()
332 26
    assert rest6.index1 == [13, 31, 49, 67, 85, 103]
333 26
    assert rest6.index2 == [109, 113, 115, 119]
334 26
    assert rest6.index3 is None
335 26
    assert rest6.index4 is None
336 26
    assert np.allclose(
337
        rest6.phase["attach"]["force_constants"], np.array([0.0, 1.25, 2.5, 3.75, 5.0])
338
    )
339 26
    assert np.allclose(
340
        rest6.phase["attach"]["targets"], np.array([0.0, 0.0, 0.0, 0.0, 0.0])
341
    )
342 26
    assert np.allclose(
343
        rest6.phase["pull"]["force_constants"], np.array([5.0, 5.0, 5.0])
344
    )
345 26
    assert np.allclose(rest6.phase["pull"]["targets"], np.array([0.0, 0.5, 1.0]))
346
    ### Note, the 6.6 in the following test is wrong ... needs to get fixed.
347 26
    assert np.allclose(
348
        rest6.phase["release"]["force_constants"], np.array([0.0, 1.65, 3.3, 4.95, 6.6])
349
    )
350 26
    assert np.allclose(
351
        rest6.phase["release"]["targets"], np.array([1.0, 1.0, 1.0, 1.0, 1.0])
352
    )
353 26
    window_list = create_window_list([rest6])
354 26
    assert window_list == [
355
        "a000",
356
        "a001",
357
        "a002",
358
        "a003",
359
        "a004",
360
        "p000",
361
        "p001",
362
        "p002",
363
        "r000",
364
        "r001",
365
        "r002",
366
        "r003",
367
        "r004",
368
    ]
369

370
    # Method 5 (Note continuous_apr = True)
371 26
    logger.info("### Testing restraint 7, Method 5")
372 26
    rest7 = DAT_restraint()
373 26
    rest7.amber_index = True
374 26
    rest7.continuous_apr = True
375 26
    rest7.auto_apr = False
376 26
    rest7.topology = os.path.join(
377
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
378
    )
379 26
    rest7.mask1 = ":1@O,O1,:BUT@H1"
380 26
    rest7.mask2 = ":CB6@N"
381 26
    rest7.attach["target"] = 0.0
382 26
    rest7.attach["fc_list"] = [0.0, 0.5, 1.0, 2.0]
383 26
    rest7.pull["fc"] = 2.0
384 26
    rest7.pull["target_list"] = [0.0, 0.5, 1.0, 1.5]
385 26
    rest7.release["target"] = 1.5
386 26
    rest7.release["fc_list"] = [0.0, 0.66, 1.2, 2.0]
387 26
    rest7.initialize()
388 26
    assert rest7.index1 == [13, 14, 111]
389 26
    assert rest7.index2 == [3]
390 26
    assert rest7.index3 is None
391 26
    assert rest7.index4 is None
392 26
    assert np.allclose(
393
        rest7.phase["attach"]["force_constants"], np.array([0.0, 0.5, 1.0, 2.0])
394
    )
395 26
    assert np.allclose(rest7.phase["attach"]["targets"], np.array([0.0, 0.0, 0.0, 0.0]))
396 26
    assert np.allclose(
397
        rest7.phase["pull"]["force_constants"], np.array([2.0, 2.0, 2.0, 2.0])
398
    )
399 26
    assert np.allclose(rest7.phase["pull"]["targets"], np.array([0.0, 0.5, 1.0, 1.5]))
400 26
    assert np.allclose(
401
        rest7.phase["release"]["force_constants"], np.array([0.0, 0.66, 1.2, 2.0])
402
    )
403 26
    assert np.allclose(
404
        rest7.phase["release"]["targets"], np.array([1.5, 1.5, 1.5, 1.5])
405
    )
406 26
    window_list = create_window_list([rest7])
407 26
    assert window_list == [
408
        "a000",
409
        "a001",
410
        "a002",
411
        "p000",
412
        "p001",
413
        "p002",
414
        "p003",
415
        "r001",
416
        "r002",
417
        "r003",
418
    ]
419

420
    # Just Attach
421 26
    logger.info("### Testing restraint 8, just attach")
422 26
    rest8 = DAT_restraint()
423 26
    rest8.amber_index = True
424 26
    rest8.continuous_apr = False
425 26
    rest8.auto_apr = False
426 26
    rest8.topology = os.path.join(
427
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
428
    )
429 26
    rest8.mask1 = ":CB6@O"
430 26
    rest8.mask2 = ":BUT@C3"
431 26
    rest8.attach["target"] = 0.0
432 26
    rest8.attach["num_windows"] = 4
433 26
    rest8.attach["fc_initial"] = 0.0
434 26
    rest8.attach["fc_final"] = 3.0
435 26
    rest8.initialize()
436 26
    assert rest8.index1 == [13]
437 26
    assert rest8.index2 == [119]
438 26
    assert rest8.index3 is None
439 26
    assert rest8.index4 is None
440 26
    assert np.allclose(
441
        rest8.phase["attach"]["force_constants"], np.array([0.0, 1.0, 2.0, 3.0])
442
    )
443 26
    assert np.allclose(rest8.phase["attach"]["targets"], np.array([0.0, 0.0, 0.0, 0.0]))
444 26
    assert rest8.phase["pull"]["force_constants"] is None
445 26
    assert rest8.phase["pull"]["targets"] is None
446 26
    assert rest8.phase["release"]["force_constants"] is None
447 26
    assert rest8.phase["release"]["targets"] is None
448 26
    window_list = create_window_list([rest8])
449 26
    assert window_list == ["a000", "a001", "a002", "a003"]
450

451
    # Just Pull
452 26
    logger.info("### Testing restraint 9, just pull")
453 26
    rest9 = DAT_restraint()
454 26
    rest9.amber_index = True
455 26
    rest9.continuous_apr = False
456 26
    rest9.auto_apr = False
457 26
    rest9.topology = os.path.join(
458
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
459
    )
460 26
    rest9.mask1 = ":CB6@O"
461 26
    rest9.mask2 = ":BUT@C3"
462 26
    rest9.pull["fc"] = 3.0
463 26
    rest9.pull["num_windows"] = 4
464 26
    rest9.pull["target_initial"] = 0.0
465 26
    rest9.pull["target_final"] = 3.0
466 26
    rest9.initialize()
467 26
    assert rest9.index1 == [13]
468 26
    assert rest9.index2 == [119]
469 26
    assert rest9.index3 is None
470 26
    assert rest9.index4 is None
471 26
    assert rest9.phase["attach"]["force_constants"] is None
472 26
    assert rest9.phase["attach"]["targets"] is None
473 26
    assert np.allclose(
474
        rest9.phase["pull"]["force_constants"], np.array([3.0, 3.0, 3.0, 3.0])
475
    )
476 26
    assert np.allclose(rest9.phase["pull"]["targets"], np.array([0.0, 1.0, 2.0, 3.0]))
477 26
    assert rest9.phase["release"]["force_constants"] is None
478 26
    assert rest9.phase["release"]["targets"] is None
479 26
    window_list = create_window_list([rest9])
480 26
    assert window_list == ["p000", "p001", "p002", "p003"]
481

482
    # Just Release
483 26
    logger.info("### Testing restraint 10, just release")
484 26
    rest10 = DAT_restraint()
485 26
    rest10.amber_index = True
486 26
    rest10.continuous_apr = False
487 26
    rest10.auto_apr = False
488 26
    rest10.topology = os.path.join(
489
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-notcentered.pdb"
490
    )
491 26
    rest10.mask1 = ":CB6@O"
492 26
    rest10.mask2 = ":BUT@C3"
493 26
    rest10.release["target"] = 0.0
494 26
    rest10.release["num_windows"] = 3
495 26
    rest10.release["fc_initial"] = 0.0
496 26
    rest10.release["fc_final"] = 2.0
497 26
    rest10.initialize()
498 26
    assert rest10.index1 == [13]
499 26
    assert rest10.index2 == [119]
500 26
    assert rest10.index3 is None
501 26
    assert rest10.index4 is None
502 26
    assert rest10.phase["attach"]["force_constants"] is None
503 26
    assert rest10.phase["attach"]["targets"] is None
504 26
    assert rest10.phase["pull"]["force_constants"] is None
505 26
    assert rest10.phase["pull"]["targets"] is None
506 26
    assert np.allclose(
507
        rest10.phase["release"]["force_constants"], np.array([0.0, 1.0, 2.0])
508
    )
509 26
    assert np.allclose(rest10.phase["release"]["targets"], np.array([0.0, 0.0, 0.0]))
510 26
    window_list = create_window_list([rest10])
511 26
    assert window_list == ["r000", "r001", "r002"]
512

513
    # Test inconsistent continuous_apr:
514 26
    with pytest.raises(Exception):
515 26
        window_list = create_window_list([rest7, rest8])
516

517
    # Test inconsistent windows:
518 26
    with pytest.raises(Exception):
519 26
        window_list = create_window_list([rest1, rest10])
520

521

522 26
def test_get_restraint_values():
523
    # Test Harmonic restraint
524 26
    attach_fractions = np.linspace(0, 1.0, 25)
525 26
    initial_distance = 2.65
526 26
    pull_distances = np.linspace(0 + initial_distance, 16.0 + initial_distance, 40)
527

528 26
    restraint = DAT_restraint()
529 26
    restraint.continuous_apr = True
530 26
    restraint.amber_index = True
531 26
    restraint.topology = os.path.join(
532
        os.path.dirname(__file__), "../data/k-cl/k-cl.pdb"
533
    )
534 26
    restraint.mask1 = "@K+"
535 26
    restraint.mask2 = "@Cl-"
536

537 26
    restraint.attach["target"] = initial_distance
538 26
    restraint.attach["fraction_list"] = attach_fractions
539 26
    restraint.attach["fc_final"] = 10.0
540

541 26
    restraint.pull["fc"] = restraint.attach["fc_final"]
542 26
    restraint.pull["target_list"] = pull_distances
543

544 26
    restraint.initialize()
545

546 26
    restraint_values = get_restraint_values(restraint, "attach", 0)
547 26
    assert restraint_values["r1"] == 0.0
548 26
    assert restraint_values["r2"] == 2.65
549 26
    assert restraint_values["r3"] == 2.65
550 26
    assert restraint_values["r4"] == 999.0
551 26
    assert restraint_values["rk2"] == 0.0
552 26
    assert restraint_values["rk3"] == 0.0
553

554 26
    restraint_values = get_restraint_values(restraint, "pull", 0)
555 26
    assert restraint_values["r1"] == 0.0
556 26
    assert restraint_values["r2"] == 2.65
557 26
    assert restraint_values["r3"] == 2.65
558 26
    assert restraint_values["r4"] == 999.0
559 26
    assert restraint_values["rk2"] == 10.0
560 26
    assert restraint_values["rk3"] == 10.0
561

562
    # Test custom values
563 26
    wall = DAT_restraint()
564 26
    wall.auto_apr = False
565 26
    wall.amber_index = True
566 26
    wall.topology = os.path.join(os.path.dirname(__file__), "../data/k-cl/k-cl.pdb")
567 26
    wall.mask1 = "@K+"
568 26
    wall.mask2 = "@Cl-"
569

570 26
    wall.attach["fc_initial"] = 1.0
571 26
    wall.attach["fc_final"] = 1.0
572

573 26
    wall.custom_restraint_values["rk2"] = 1.0
574 26
    wall.custom_restraint_values["rk3"] = 1.0
575 26
    wall.custom_restraint_values["r2"] = 0.0
576 26
    wall.custom_restraint_values["r3"] = 3.5
577

578 26
    wall.attach["target"] = 3.5
579 26
    wall.attach["num_windows"] = len(attach_fractions)
580

581 26
    wall.initialize()
582

583 26
    restraint_values = get_restraint_values(wall, "attach", 0)
584 26
    assert restraint_values["r1"] == 0.0
585 26
    assert restraint_values["r2"] == 0.0
586 26
    assert restraint_values["r3"] == 3.5
587 26
    assert restraint_values["r4"] == 999.0
588 26
    assert restraint_values["rk2"] == 1.0
589 26
    assert restraint_values["rk3"] == 1.0
590

591 26
    wall = DAT_restraint()
592 26
    wall.auto_apr = False
593 26
    wall.amber_index = True
594 26
    wall.topology = os.path.join(os.path.dirname(__file__), "../data/k-cl/k-cl.pdb")
595 26
    wall.mask1 = "@K+"
596 26
    wall.mask2 = "@Cl-"
597

598 26
    wall.attach["fc_initial"] = 1.0
599 26
    wall.attach["fc_final"] = 1.0
600

601 26
    wall.custom_restraint_values["rk2"] = 1.0
602 26
    wall.custom_restraint_values["rk3"] = 1.0
603 26
    wall.custom_restraint_values["r2"] = 3.5
604 26
    wall.custom_restraint_values["r3"] = 0.0
605

606 26
    wall.attach["target"] = 3.5
607 26
    wall.attach["num_windows"] = len(attach_fractions)
608

609 26
    wall.initialize()
610

611 26
    restraint_values = get_restraint_values(wall, "attach", 0)
612 26
    assert restraint_values["r1"] == 0.0
613 26
    assert restraint_values["r2"] == 3.5
614 26
    assert restraint_values["r3"] == 0.0
615 26
    assert restraint_values["r4"] == 999.0
616 26
    assert restraint_values["rk2"] == 1.0
617 26
    assert restraint_values["rk3"] == 1.0
618

619

620 26
def test_get_bias_potential_type():
621
    # Test Harmonic restraint
622 26
    attach_fractions = np.linspace(0, 1.0, 25)
623 26
    initial_distance = 2.65
624 26
    pull_distances = np.linspace(0 + initial_distance, 16.0 + initial_distance, 40)
625

626 26
    restraint = DAT_restraint()
627 26
    restraint.continuous_apr = True
628 26
    restraint.amber_index = True
629 26
    restraint.topology = os.path.join(
630
        os.path.dirname(__file__), "../data/k-cl/k-cl.pdb"
631
    )
632 26
    restraint.mask1 = "@K+"
633 26
    restraint.mask2 = "@Cl-"
634

635 26
    restraint.attach["target"] = initial_distance
636 26
    restraint.attach["fraction_list"] = attach_fractions
637 26
    restraint.attach["fc_final"] = 10.0
638

639 26
    restraint.pull["fc"] = restraint.attach["fc_final"]
640 26
    restraint.pull["target_list"] = pull_distances
641

642 26
    restraint.initialize()
643

644 26
    assert get_bias_potential_type(restraint, "attach", 0) == "restraint"
645 26
    assert get_bias_potential_type(restraint, "pull", 0) == "restraint"
646

647
    # Test upper wall restraint (1)
648 26
    upper = DAT_restraint()
649 26
    upper.auto_apr = False
650 26
    upper.amber_index = True
651 26
    upper.topology = os.path.join(os.path.dirname(__file__), "../data/k-cl/k-cl.pdb")
652 26
    upper.mask1 = "@K+"
653 26
    upper.mask2 = "@Cl-"
654

655 26
    upper.attach["fc_initial"] = 1.0
656 26
    upper.attach["fc_final"] = 1.0
657

658 26
    upper.custom_restraint_values["rk2"] = 1.0
659 26
    upper.custom_restraint_values["rk3"] = 1.0
660 26
    upper.custom_restraint_values["r2"] = 0.0
661 26
    upper.custom_restraint_values["r3"] = 3.5
662

663 26
    upper.attach["target"] = 3.5
664 26
    upper.attach["num_windows"] = len(attach_fractions)
665

666 26
    upper.initialize()
667

668 26
    assert get_bias_potential_type(upper, "attach", 0) == "upper_walls"
669 26
    assert get_bias_potential_type(upper, "attach", 1) == "upper_walls"
670

671
    # Test upper wall restraint (2)
672 26
    upper = DAT_restraint()
673 26
    upper.auto_apr = False
674 26
    upper.amber_index = True
675 26
    upper.topology = os.path.join(os.path.dirname(__file__), "../data/k-cl/k-cl.pdb")
676 26
    upper.mask1 = "@K+"
677 26
    upper.mask2 = "@Cl-"
678

679 26
    upper.attach["fc_initial"] = 1.0
680 26
    upper.attach["fc_final"] = 1.0
681

682 26
    upper.custom_restraint_values["rk2"] = 0.0
683 26
    upper.custom_restraint_values["rk3"] = 1.0
684 26
    upper.custom_restraint_values["r2"] = 3.5
685 26
    upper.custom_restraint_values["r3"] = 3.5
686

687 26
    upper.attach["target"] = 3.5
688 26
    upper.attach["num_windows"] = len(attach_fractions)
689

690 26
    upper.initialize()
691

692 26
    assert get_bias_potential_type(upper, "attach", 0) == "upper_walls"
693 26
    assert get_bias_potential_type(upper, "attach", 1) == "upper_walls"
694

695
    # Test upper wall restraint (3)
696 26
    upper = DAT_restraint()
697 26
    upper.auto_apr = False
698 26
    upper.amber_index = True
699 26
    upper.topology = os.path.join(os.path.dirname(__file__), "../data/k-cl/k-cl.pdb")
700 26
    upper.mask1 = "@K+"
701 26
    upper.mask2 = "@Cl-"
702

703 26
    upper.attach["fc_initial"] = 1.0
704 26
    upper.attach["fc_final"] = 1.0
705

706 26
    upper.custom_restraint_values["rk2"] = 0.0
707 26
    upper.custom_restraint_values["rk3"] = 1.0
708 26
    upper.custom_restraint_values["r2"] = 0.0
709 26
    upper.custom_restraint_values["r3"] = 3.5
710

711 26
    upper.attach["target"] = 3.5
712 26
    upper.attach["num_windows"] = len(attach_fractions)
713

714 26
    upper.initialize()
715

716 26
    assert get_bias_potential_type(upper, "attach", 0) == "upper_walls"
717 26
    assert get_bias_potential_type(upper, "attach", 1) == "upper_walls"
718

719
    # Test lower wall restraint (1)
720 26
    lower = DAT_restraint()
721 26
    lower.auto_apr = False
722 26
    lower.amber_index = True
723 26
    lower.topology = os.path.join(os.path.dirname(__file__), "../data/k-cl/k-cl.pdb")
724 26
    lower.mask1 = "@K+"
725 26
    lower.mask2 = "@Cl-"
726

727 26
    lower.attach["fc_initial"] = 1.0
728 26
    lower.attach["fc_final"] = 1.0
729

730 26
    lower.custom_restraint_values["rk2"] = 1.0
731 26
    lower.custom_restraint_values["rk3"] = 1.0
732 26
    lower.custom_restraint_values["r2"] = 3.5
733 26
    lower.custom_restraint_values["r3"] = 0.0
734

735 26
    lower.attach["target"] = 3.5
736 26
    lower.attach["num_windows"] = len(attach_fractions)
737

738 26
    lower.initialize()
739

740 26
    assert get_bias_potential_type(lower, "attach", 0) == "lower_walls"
741 26
    assert get_bias_potential_type(lower, "attach", 1) == "lower_walls"
742

743
    # Test lower wall restraint (2)
744 26
    lower = DAT_restraint()
745 26
    lower.auto_apr = False
746 26
    lower.amber_index = True
747 26
    lower.topology = os.path.join(os.path.dirname(__file__), "../data/k-cl/k-cl.pdb")
748 26
    lower.mask1 = "@K+"
749 26
    lower.mask2 = "@Cl-"
750

751 26
    lower.attach["fc_initial"] = 1.0
752 26
    lower.attach["fc_final"] = 1.0
753

754 26
    lower.custom_restraint_values["rk2"] = 1.0
755 26
    lower.custom_restraint_values["rk3"] = 0.0
756 26
    lower.custom_restraint_values["r2"] = 3.5
757 26
    lower.custom_restraint_values["r3"] = 3.5
758

759 26
    lower.attach["target"] = 3.5
760 26
    lower.attach["num_windows"] = len(attach_fractions)
761

762 26
    lower.initialize()
763

764 26
    assert get_bias_potential_type(lower, "attach", 0) == "lower_walls"
765 26
    assert get_bias_potential_type(lower, "attach", 1) == "lower_walls"
766

767
    # Test lower wall restraint (3)
768 26
    lower = DAT_restraint()
769 26
    lower.auto_apr = False
770 26
    lower.amber_index = True
771 26
    lower.topology = os.path.join(os.path.dirname(__file__), "../data/k-cl/k-cl.pdb")
772 26
    lower.mask1 = "@K+"
773 26
    lower.mask2 = "@Cl-"
774

775 26
    lower.attach["fc_initial"] = 1.0
776 26
    lower.attach["fc_final"] = 1.0
777

778 26
    lower.custom_restraint_values["rk2"] = 1.0
779 26
    lower.custom_restraint_values["rk3"] = 0.0
780 26
    lower.custom_restraint_values["r2"] = 3.5
781 26
    lower.custom_restraint_values["r3"] = 6.5
782

783 26
    lower.attach["target"] = 3.5
784 26
    lower.attach["num_windows"] = len(attach_fractions)
785

786 26
    lower.initialize()
787

788 26
    assert get_bias_potential_type(lower, "attach", 0) == "lower_walls"
789 26
    assert get_bias_potential_type(lower, "attach", 1) == "lower_walls"
790

791

792 26
def test_extract_guest_restraints():
793
    """Test extract_guest_restraints to make sure it is extracting the correct restraints."""
794 26
    restraints = []
795

796
    # Guest - r
797 26
    r = DAT_restraint()
798 26
    r.amber_index = True
799 26
    r.continuous_apr = True
800 26
    r.auto_apr = True
801 26
    r.topology = os.path.join(
802
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-dum.pdb"
803
    )
804 26
    r.mask1 = ":DM1"
805 26
    r.mask2 = ":BUT@C3"
806 26
    r.attach["target"] = 6.0
807 26
    r.attach["num_windows"] = 4
808 26
    r.attach["fc_initial"] = 0.0
809 26
    r.attach["fc_final"] = 5.0
810 26
    r.pull["fc"] = r.attach["fc_final"]
811 26
    r.pull["num_windows"] = 46
812 26
    r.pull["target_initial"] = r.attach["target"]
813 26
    r.pull["target_final"] = 24.0
814 26
    r.release["target"] = r.pull["target_final"]
815 26
    r.release["num_windows"] = r.attach["num_windows"]
816 26
    r.release["fc_initial"] = r.attach["fc_initial"]
817 26
    r.release["fc_final"] = r.attach["fc_final"]
818 26
    r.initialize()
819 26
    restraints.append(r)
820

821
    # Guest - theta
822 26
    r = DAT_restraint()
823 26
    r.amber_index = True
824 26
    r.continuous_apr = True
825 26
    r.auto_apr = True
826 26
    r.topology = os.path.join(
827
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-dum.pdb"
828
    )
829 26
    r.mask1 = ":DM1"
830 26
    r.mask2 = ":BUT@C3"
831 26
    r.mask3 = ":BUT@C"
832 26
    r.attach["target"] = 180.0
833 26
    r.attach["num_windows"] = 15
834 26
    r.attach["fc_initial"] = 0.0
835 26
    r.attach["fc_final"] = 100.0
836 26
    r.pull["fc"] = r.attach["fc_final"]
837 26
    r.pull["num_windows"] = 46
838 26
    r.pull["target_initial"] = r.attach["target"]
839 26
    r.pull["target_final"] = 180.0
840 26
    r.release["target"] = r.pull["target_final"]
841 26
    r.release["num_windows"] = r.attach["num_windows"]
842 26
    r.release["fc_initial"] = r.attach["fc_initial"]
843 26
    r.release["fc_final"] = r.attach["fc_final"]
844 26
    r.initialize()
845 26
    restraints.append(r)
846

847
    # Guest - beta
848 26
    r = DAT_restraint()
849 26
    r.amber_index = True
850 26
    r.continuous_apr = True
851 26
    r.auto_apr = True
852 26
    r.topology = os.path.join(
853
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-dum.pdb"
854
    )
855 26
    r.mask1 = ":DM2"
856 26
    r.mask2 = ":DM1"
857 26
    r.mask3 = ":BUT@C3"
858 26
    r.attach["target"] = 180.0
859 26
    r.attach["num_windows"] = 15
860 26
    r.attach["fc_initial"] = 0.0
861 26
    r.attach["fc_final"] = 100.0
862 26
    r.pull["fc"] = r.attach["fc_final"]
863 26
    r.pull["num_windows"] = 46
864 26
    r.pull["target_initial"] = r.attach["target"]
865 26
    r.pull["target_final"] = 180.0
866 26
    r.release["target"] = r.pull["target_final"]
867 26
    r.release["num_windows"] = r.attach["num_windows"]
868 26
    r.release["fc_initial"] = r.attach["fc_initial"]
869 26
    r.release["fc_final"] = r.attach["fc_final"]
870 26
    r.initialize()
871 26
    restraints.append(r)
872

873 26
    structure = pmd.load_file(
874
        os.path.join(os.path.dirname(__file__), "../data/cb6-but/cb6-but-dum.pdb"),
875
        structure=True,
876
    )
877 26
    guest_restraints = extract_guest_restraints(structure, restraints, "BUT")
878

879 26
    assert guest_restraints[0] is not None
880 26
    assert guest_restraints[1] is not None
881 26
    assert guest_restraints[2] is None
882 26
    assert guest_restraints[3] is None
883 26
    assert guest_restraints[4] is not None
884 26
    assert guest_restraints[5] is None
885

886
    # Guest - phi
887 26
    r = DAT_restraint()
888 26
    r.amber_index = True
889 26
    r.continuous_apr = True
890 26
    r.auto_apr = True
891 26
    r.topology = os.path.join(
892
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-dum.pdb"
893
    )
894 26
    r.mask1 = ":DM3"
895 26
    r.mask2 = ":DM2"
896 26
    r.mask3 = ":DM1"
897 26
    r.mask4 = ":BUT@C3"
898 26
    r.attach["target"] = 120.0
899 26
    r.attach["num_windows"] = 15
900 26
    r.attach["fc_initial"] = 0.0
901 26
    r.attach["fc_final"] = 100.0
902 26
    r.pull["fc"] = r.attach["fc_final"]
903 26
    r.pull["num_windows"] = 46
904 26
    r.pull["target_initial"] = r.attach["target"]
905 26
    r.pull["target_final"] = 120.0
906 26
    r.release["target"] = r.pull["target_final"]
907 26
    r.release["num_windows"] = r.attach["num_windows"]
908 26
    r.release["fc_initial"] = r.attach["fc_initial"]
909 26
    r.release["fc_final"] = r.attach["fc_final"]
910 26
    r.initialize()
911 26
    restraints.append(r)
912

913 26
    guest_restraints = extract_guest_restraints(structure, restraints, "BUT")
914

915 26
    assert guest_restraints[0] is not None
916 26
    assert guest_restraints[1] is not None
917 26
    assert guest_restraints[2] is not None
918 26
    assert guest_restraints[3] is None
919 26
    assert guest_restraints[4] is not None
920 26
    assert guest_restraints[5] is None
921

922
    # Guest - alpha
923 26
    r = DAT_restraint()
924 26
    r.amber_index = True
925 26
    r.continuous_apr = True
926 26
    r.auto_apr = True
927 26
    r.topology = os.path.join(
928
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-dum.pdb"
929
    )
930 26
    r.mask1 = ":DM2"
931 26
    r.mask2 = ":DM1"
932 26
    r.mask3 = ":BUT@C3"
933 26
    r.mask4 = ":BUT@C"
934 26
    r.attach["target"] = 50.0
935 26
    r.attach["num_windows"] = 15
936 26
    r.attach["fc_initial"] = 0.0
937 26
    r.attach["fc_final"] = 100.0
938 26
    r.pull["fc"] = r.attach["fc_final"]
939 26
    r.pull["num_windows"] = 46
940 26
    r.pull["target_initial"] = r.attach["target"]
941 26
    r.pull["target_final"] = 50.0
942 26
    r.release["target"] = r.pull["target_final"]
943 26
    r.release["num_windows"] = r.attach["num_windows"]
944 26
    r.release["fc_initial"] = r.attach["fc_initial"]
945 26
    r.release["fc_final"] = r.attach["fc_final"]
946 26
    r.initialize()
947 26
    restraints.append(r)
948

949 26
    guest_restraints = extract_guest_restraints(structure, restraints, "BUT")
950

951 26
    assert guest_restraints[0] is not None
952 26
    assert guest_restraints[1] is not None
953 26
    assert guest_restraints[2] is not None
954 26
    assert guest_restraints[3] is not None
955 26
    assert guest_restraints[4] is not None
956 26
    assert guest_restraints[5] is None
957

958
    # Guest - gamma
959 26
    r = DAT_restraint()
960 26
    r.amber_index = True
961 26
    r.continuous_apr = True
962 26
    r.auto_apr = True
963 26
    r.topology = os.path.join(
964
        os.path.dirname(__file__), "../data/cb6-but/cb6-but-dum.pdb"
965
    )
966 26
    r.mask1 = ":DM1"
967 26
    r.mask2 = ":BUT@C3"
968 26
    r.mask3 = ":BUT@C"
969 26
    r.mask4 = ":BUT@C2"
970 26
    r.attach["target"] = 50.0
971 26
    r.attach["num_windows"] = 15
972 26
    r.attach["fc_initial"] = 0.0
973 26
    r.attach["fc_final"] = 100.0
974 26
    r.pull["fc"] = r.attach["fc_final"]
975 26
    r.pull["num_windows"] = 46
976 26
    r.pull["target_initial"] = r.attach["target"]
977 26
    r.pull["target_final"] = 50.0
978 26
    r.release["target"] = r.pull["target_final"]
979 26
    r.release["num_windows"] = r.attach["num_windows"]
980 26
    r.release["fc_initial"] = r.attach["fc_initial"]
981 26
    r.release["fc_final"] = r.attach["fc_final"]
982 26
    r.initialize()
983 26
    restraints.append(r)
984

985 26
    guest_restraints = extract_guest_restraints(structure, restraints, "BUT")
986

987 26
    assert guest_restraints[0] is not None
988 26
    assert guest_restraints[1] is not None
989 26
    assert guest_restraints[2] is not None
990 26
    assert guest_restraints[3] is not None
991 26
    assert guest_restraints[4] is not None
992 26
    assert guest_restraints[5] is not None

Read our documentation on viewing source code .

Loading