1 2
import logging
2

3 2
import numpy as np
4 2
import parmed as pmd
5 2
import pytraj as pt
6 2
from paprika import utils
7

8 2
logger = logging.getLogger(__name__)
9

10

11 2
class DAT_restraint(object):
12
    """
13
    Distance or angle or torsion restraints on atoms in the simulation.
14
    """
15

16 2
    @property
17
    def instances(self):
18
        """A list of ``DAT_restraints`` that have been initialized.
19

20
        .. note ::
21

22
        This should never be called directly and ought to be private.
23
        """
24 0
        return self._instances
25

26 2
    @instances.setter
27
    def instances(self, value):
28 0
        self._instances = value
29

30 2
    @property
31
    def topology(self):
32
        """The topology file used to initialize the restraints. This is often a PDB file."""
33 2
        return self._topology
34

35 2
    @topology.setter
36
    def topology(self, value):
37 2
        self._topology = value
38

39 2
    @property
40
    def mask1(self):
41
        """The first atom mask used for the restraint."""
42 2
        return self._mask1
43

44 2
    @mask1.setter
45
    def mask1(self, value):
46 2
        self._mask1 = value
47

48 2
    @property
49
    def mask2(self):
50
        """The second atom mask used for the restraint."""
51 2
        return self._mask2
52

53 2
    @mask2.setter
54
    def mask2(self, value):
55 2
        self._mask2 = value
56

57 2
    @property
58
    def mask3(self):
59
        """The third atom mask used for the restraint. Only two atom masks are required."""
60 2
        return self._mask3
61

62 2
    @mask3.setter
63
    def mask3(self, value):
64 2
        self._mask3 = value
65

66 2
    @property
67
    def mask4(self):
68
        """The fourth atom mask used for the restraint. Only two atom masks are required."""
69 2
        return self._mask4
70

71 2
    @mask4.setter
72
    def mask4(self, value):
73 2
        self._mask4 = value
74

75 2
    @property
76
    def custom_restraint_values(self):
77
        """In the case a non-harmonic restraint is desired, the pre-calculated values (r1, r2, and so on) can be
78
        overridden with ones from this dictionary. These values will be directly written to the AMBER restraint file,
79
        so the keys must be valid AMBER keywords.
80

81
        Specifically, for target distance, R:
82

83
        • R < r1 Linear, with the slope of the "left-hand" parabola at the point R = r1.
84
        • r1 <= R < r2 Parabolic, with restraint energy k2(R − r2)².
85
        • r2 <= R < r3 E = 0.
86
        • r3 <= R < r4 Parabolic, with restraint energy k3(R − r3)².
87
        • r4 <= R Linear, with the slope of the "right-hand" parabola at the point R = r4.
88

89
        In the case of AMBER18, this is covered in section 25.1 of the manual.
90

91
        For example, a flat bottom restraint can be specified by setting ``r2`` not equal to ``r3``.
92

93
        >>> from paprika.restraints import DAT_restraint
94
        >>> this = DAT_restraint
95
        >>> this.custom_restraint_values["r2"] = 0.0
96
        >>> this.custom_restraint_values["r3"] = 12.0
97

98
        The values in this dictionary use default AMBER units for distances and force constants.
99

100
        This attribute has **no effect** for OpenMM.
101
        """
102 2
        return self._custom_restraint_values
103

104 2
    @custom_restraint_values.setter
105
    def custom_restraint_values(self, value):
106 0
        self._custom_restraint_values = value
107

108 2
    @property
109
    def auto_apr(self):
110
        """If ``True``, (1) the force constant during the pulling phase will be set to the final force constant from
111
        the attach phase, and (2) the initial target during the pulling phase will be equal to the final target from
112
        the attach phase. If ``False``, these values must be set manually.
113
        """
114 2
        return self._auto_apr
115

116 2
    @auto_apr.setter
117
    def auto_apr(self, value):
118 2
        if not isinstance(value, bool):
119 0
            raise TypeError("The ``auto_apr`` attribute is a boolean.")
120 2
        self._auto_apr = value
121

122 2
    @property
123
    def continuous_apr(self):
124
        """
125
        If ``True``, (1) the final window of the attach phase is used as the first window of the pull phase, and (2) the
126
        final window of the pull phase is used as the first window of the release phase.
127
        """
128 2
        return self._continuous_apr
129

130 2
    @continuous_apr.setter
131
    def continuous_apr(self, value):
132 2
        if not isinstance(value, bool):
133 0
            raise TypeError("The ``continuous_apr`` attribute is a boolean.")
134 2
        self._continuous_apr = value
135

136 2
    @property
137
    def attach(self):
138
        """
139
        Dictionary specifying the APR parameters during the attach phase. The dictionary keys are as follows:
140

141
            - ``target``             : The target value for the restraint (mandatory)
142
            - ``fc_initial``         : The initial force constant (optional)
143
            - ``fc_final``           : The final force constant (optional)
144
            - ``num_windows``        : The number of windows (optional)
145
            - ``fc_increment``       : The force constant increment (optional)
146
            - ``fraction_increment`` : The percentage of the force constant increment (optional)
147
            - ``fraction_list``      : The list of force constant percentages (optional)
148
            - ``fc_list``            : The list of force constants (will be created if not given)
149

150
        .. note ::
151
            This is fragile and this could be hardened by making these ``ENUM``s and doing much more type-checking.
152
        """
153

154 2
        return self._attach
155

156 2
    @attach.setter
157
    def attach(self, value):
158 0
        self._attach = value
159

160 2
    @property
161
    def pull(self):
162
        """
163
        Dictionary specifying the APR parameters during the pull phase. The dictionary keys are as follows:
164

165
            - ``fc``                 : The force constant for the restraint (mandatory)
166
            - ``target_initial``     : The initial target value (optional)
167
            - ``target_final``       : The final target value (optional)
168
            - ``num_windows``        : The number of windows (optional)
169
            - ``target_increment``   : The target value increment (optional)
170
            - ``fraction_increment`` : The percentage of the target value increment (optional)
171
            - ``fraction_list``      : The list of target value percentages (optional)
172
            - ``target_list``        : The list of target values (will be created if not given)
173

174
        .. note ::
175
            This is fragile and this could be hardened by making these ``ENUM``s and doing much more type-checking.
176
        """
177

178 2
        return self._pull
179

180 2
    @pull.setter
181
    def pull(self, value):
182 0
        self._pull = value
183

184 2
    @property
185
    def release(self):
186
        """
187
        Dictionary specifying the APR parameters during the release phase. The dictionary keys are as follows:
188

189
            - ``target``             : The target value for the restraint (mandatory)
190
            - ``fc_initial``         : The initial force constant (optional)
191
            - ``fc_final``           : The final force constant (optional)
192
            - ``num_windows``        : The number of windows (optional)
193
            - ``fc_increment``       : The force constant increment (optional)
194
            - ``fraction_increment`` : The percentage of the force constant increment (optional)
195
            - ``fraction_list``      : The list of force constant percentages (optional)
196
            - ``fc_list``            : The list of force constants (will be created if not given)
197

198
        .. note ::
199
            This is fragile and this could be hardened by making these ``ENUM``s and doing much more type-checking.
200
        """
201 2
        return self._release
202

203 2
    @release.setter
204
    def release(self, value):
205 0
        self._release = value
206

207 2
    @property
208
    def amber_index(self):
209
        """
210
        If ``True``, add 1 to all atom indices.
211
        """
212 2
        return self._amber_index
213

214 2
    @amber_index.setter
215
    def amber_index(self, value):
216 2
        self._amber_index = value
217

218 2
    instances = []
219

220 2
    def __init__(self):
221

222 2
        self._topology = None
223 2
        self._mask1 = None
224 2
        self._mask2 = None
225 2
        self._mask3 = None
226 2
        self._mask4 = None
227

228
        # These indices will be automatically populated during :meth:`paprika.restraints.DAT_restraint.initialize`.
229 2
        self.index1 = None
230 2
        self.index2 = None
231 2
        self.index3 = None
232 2
        self.index4 = None
233

234 2
        self.group1 = False
235 2
        self.group2 = False
236 2
        self.group3 = False
237 2
        self.group4 = False
238

239 2
        self._custom_restraint_values = {}
240

241 2
        self._auto_apr = False
242 2
        self._continuous_apr = True
243

244 2
        self._amber_index = False
245

246 2
        self._attach = {
247
            "target": None,
248
            "fc_initial": None,
249
            "fc_final": None,
250
            "num_windows": None,
251
            "fc_increment": None,
252
            "fraction_increment": None,
253
            "fraction_list": None,
254
            "fc_list": None,
255
        }
256 2
        self._pull = {
257
            "fc": None,
258
            "target_initial": None,
259
            "target_final": None,
260
            "num_windows": None,
261
            "target_increment": None,
262
            "fraction_increment": None,
263
            "fraction_list": None,
264
            "target_list": None,
265
        }
266 2
        self._release = {
267
            "target": None,
268
            "fc_initial": None,
269
            "fc_final": None,
270
            "num_windows": None,
271
            "fc_increment": None,
272
            "fraction_increment": None,
273
            "fraction_list": None,
274
            "fc_list": None,
275
        }
276

277 2
        DAT_restraint.instances.append(self)
278

279 2
    def __eq__(self, other):
280
        """
281
        Test whether two ``DAT_restraint`` instances are equivalent.
282
        """
283 2
        self_dictionary = self.__dict__
284 2
        other_dictionary = other.__dict__
285 2
        for dct in [self_dictionary, other_dictionary]:
286
            # Skip checking topology.
287
            # It is difficult to check topology for two reasons.
288
            # One, there are numerical variances in the numeric parameters.
289
            # We can use `np.allcose()` but it needs to be done manually on the
290
            # nested ParmEd AmberParm or Structure classes.
291
            # Two, sometimes the dummy atoms seem to be encoded with index `-1`.
292
            # This can happen if the structure used to setup the restraint
293
            # does not have the dummy atoms already included.
294
            # Previously, I checked if the topology was an AmberParm, and tried
295
            # to replace that with the string representation of the file name,
296
            # but for some reason, that attribute was not always available.
297
            # Thus, I simply do not compare the `topology` attribute.
298 2
            dct["topology"] = None
299 2
        logger.debug(self_dictionary)
300 2
        logger.debug(other_dictionary)
301 2
        keys = set(self_dictionary.keys()) & set(other_dictionary.keys())
302 2
        for key in keys:
303 2
            if key != "phase":
304 2
                assert self_dictionary[key] == other_dictionary[key]
305
            else:
306 2
                for phs in ["attach", "pull", "release"]:
307 2
                    for value in ["force_constants", "targets"]:
308 2
                        if (
309
                            self_dictionary["phase"][phs][value] is None
310
                            and other_dictionary["phase"][phs][value] is None
311
                        ):
312 0
                            continue
313
                        else:
314 2
                            assert np.allclose(
315
                                self_dictionary["phase"][phs][value],
316
                                other_dictionary["phase"][phs][value],
317
                            )
318 2
        return True
319

320 2
    @staticmethod
321
    def _calc_method(phase, restraint_dictionary, method):
322
        """
323
        This helper function figures out which values in the restraint dictionary need to be set.
324

325
        phase: str
326
            The restraint phase.
327
        restraint_dictionary: dict
328
            The restraint dictionary corresponding whose values will be set.
329
        method: str
330
            Which "method" will be used to set the remaining values. The "method" refers to which
331
            combination of restraint inputs are provided -- for example, initial value, final value,
332
            and number of windows or initial value, increment size, and number of increments -- which
333
            are defined in :meth:`paprika.restraints.DAT_restraint.initialize`.
334

335
            This appears to be a ``str`` but should be an ``int``.
336

337
        .. note ::
338
            This could do with some serious sprucing up.
339

340
        """
341

342 2
        force_constants = None
343 2
        targets = None
344

345
        # Attach/Release, Force Constant Method 1
346 2
        if phase in ("a", "r") and method == "1":
347 2
            force_constants = np.linspace(
348
                restraint_dictionary["fc_initial"],
349
                restraint_dictionary["fc_final"],
350
                restraint_dictionary["num_windows"],
351
            )
352

353
        # Attach/Release, Force Constant Method 1a
354 2
        elif phase in ("a", "r") and method == "1a":
355 2
            force_constants = np.linspace(
356
                0.0,
357
                restraint_dictionary["fc_final"],
358
                restraint_dictionary["num_windows"],
359
            )
360

361
        # Attach/Release, Force Constant Method 2
362 2
        elif phase in ("a", "r") and method == "2":
363 2
            force_constants = np.arange(
364
                restraint_dictionary["fc_initial"],
365
                restraint_dictionary["fc_final"] + restraint_dictionary["fc_increment"],
366
                restraint_dictionary["fc_increment"],
367
            )
368

369
        # Attach/Release, Force Constant Method 2a
370 2
        elif phase in ("a", "r") and method == "2a":
371 2
            force_constants = np.arange(
372
                0.0,
373
                restraint_dictionary["fc_final"] + restraint_dictionary["fc_increment"],
374
                restraint_dictionary["fc_increment"],
375
            )
376

377
        # Attach/Release, Force Constant Method 3
378 2
        elif phase in ("a", "r") and method == "3":
379 2
            force_constants = np.asarray(
380
                [
381
                    fraction * restraint_dictionary["fc_final"]
382
                    for fraction in restraint_dictionary["fraction_list"]
383
                ]
384
            )
385

386
        # Attach/Release, Force Constant Method 4
387 2
        elif phase in ("a", "r") and method == "4":
388 2
            fractions = np.arange(
389
                0,
390
                1.0 + restraint_dictionary["fraction_increment"],
391
                restraint_dictionary["fraction_increment"],
392
            )
393 2
            force_constants = np.asarray(
394
                [fraction * restraint_dictionary["fc_final"] for fraction in fractions]
395
            )
396

397
        # Attach/Release, Force Constant Method 5
398 2
        elif phase in ("a", "r") and method == "5":
399 2
            force_constants = np.asarray(restraint_dictionary["fc_list"])
400

401
        # Attach/Release, Target Method
402 2
        if phase in ("a", "r"):
403 2
            targets = np.asarray(
404
                [restraint_dictionary["target"]] * len(force_constants)
405
            )
406

407
        # Pull, Target Method 1
408 2
        if phase == "p" and method == "1":
409 2
            targets = np.linspace(
410
                restraint_dictionary["target_initial"],
411
                restraint_dictionary["target_final"],
412
                restraint_dictionary["num_windows"],
413
            )
414

415
        # Pull, Target Method 1a
416 2
        elif phase == "p" and method == "1a":
417 2
            targets = np.linspace(
418
                0.0,
419
                restraint_dictionary["target_final"],
420
                restraint_dictionary["num_windows"],
421
            )
422

423
        # Pull, Target Method 2
424 2
        elif phase == "p" and method == "2":
425 2
            targets = np.arange(
426
                restraint_dictionary["target_initial"],
427
                restraint_dictionary["target_final"]
428
                + restraint_dictionary["target_increment"],
429
                restraint_dictionary["target_increment"],
430
            )
431

432
        # Pull, Target Method 2a
433 2
        elif phase == "p" and method == "2a":
434 2
            targets = np.arange(
435
                0.0,
436
                restraint_dictionary["target_final"]
437
                + restraint_dictionary["target_increment"],
438
                restraint_dictionary["target_increment"],
439
            )
440

441
        # Pull, Target Method 3
442 2
        elif phase == "p" and method == "3":
443 2
            targets = np.asarray(
444
                [
445
                    fraction * restraint_dictionary["target_final"]
446
                    for fraction in restraint_dictionary["fraction_list"]
447
                ]
448
            )
449

450
        # Pull, Target Method 4
451 2
        elif phase == "p" and method == "4":
452 2
            fractions = np.arange(
453
                0,
454
                1.0 + restraint_dictionary["fraction_increment"],
455
                restraint_dictionary["fraction_increment"],
456
            )
457 2
            targets = np.asarray(
458
                [
459
                    fraction * restraint_dictionary["target_final"]
460
                    for fraction in fractions
461
                ]
462
            )
463

464
        # Pull, Target Method 5
465 2
        elif phase == "p" and method == "5":
466 2
            targets = np.asarray(restraint_dictionary["target_list"])
467

468
        # Pull, Force Constant Method
469 2
        if phase == "p":
470 2
            force_constants = np.asarray([restraint_dictionary["fc"]] * len(targets))
471

472 2
        if force_constants is None and targets is None:
473 0
            logger.error("Unsupported Phase/Method: {} / {}".format(phase, method))
474 0
            raise Exception("Unexpected phase/method combination passed to _calc_meth")
475

476 2
        return force_constants, targets
477

478 2
    def initialize(self):
479
        """
480
        Automatically set remaining force constants and targets.
481

482
        Depending on which values are provided for each phase, a different method will
483
        be used to determine the list of force constants and targets (below).
484

485
        For attach and release, a ``target`` value is required and the method is determined if the
486
        following values are not ``None``:
487

488
            - Method 1:   num_windows, fc_initial, fc_final
489
            - Method 1a:  num_windows, fc_final
490
            - Method 2:   fc_increment, fc_initial, fc_final
491
            - Method 2a:  fc_increment, fc_final
492
            - Method 3:   fraction_list, fc_final
493
            - Method 4:   fraction_increment, fc_final
494
            - Method 5:   fc_list
495

496
        For pull, a ``fc`` value is required and the method is determined if the
497
        following values are not ``None``:
498

499
            - Method 1:   num_windows, target_initial, target_final
500
            - Method 1a:  num_windows, target_final
501
            - Method 2:   target_increment, target_initial, target_final
502
            - Method 2a:  target_increment, target_final
503
            - Method 3:   fraction_list, target_final
504
            - Method 4:   fraction_increment, target_final
505
            - Method 5:   target_list
506

507
        .. note ::
508
            This is unnecessary overengineering.
509
        """
510

511 2
        self.phase = {
512
            "attach": {"force_constants": None, "targets": None},
513
            "pull": {"force_constants": None, "targets": None},
514
            "release": {"force_constants": None, "targets": None},
515
        }
516
        # ------------------------------------ ATTACH ------------------------------------ #
517 2
        logger.debug("Calculating attach targets and force constants...")
518

519
        # Temporary variables to improve readability
520 2
        force_constants = None
521 2
        targets = None
522

523 2
        if (
524
            self.attach["num_windows"] is not None
525
            and self.attach["fc_final"] is not None
526
        ):
527 2
            if self.attach["fc_initial"] is not None:
528 2
                logger.debug("Attach, Method #1")
529 2
                force_constants, targets = self._calc_method("a", self.attach, "1")
530
            else:
531 2
                logger.debug("Attach, Method #1a")
532 2
                force_constants, targets = self._calc_method("a", self.attach, "1a")
533

534 2
        elif (
535
            self.attach["fc_increment"] is not None
536
            and self.attach["fc_final"] is not None
537
        ):
538 2
            if self.attach["fc_initial"] is not None:
539 2
                logger.debug("Attach, Method #2")
540 2
                force_constants, targets = self._calc_method("a", self.attach, "2")
541
            else:
542 2
                logger.debug("Attach, Method #2a")
543 2
                force_constants, targets = self._calc_method("a", self.attach, "2a")
544

545 2
        elif (
546
            self.attach["fraction_list"] is not None
547
            and self.attach["fc_final"] is not None
548
        ):
549 2
            logger.debug("Attach, Method #3")
550 2
            force_constants, targets = self._calc_method("a", self.attach, "3")
551

552 2
        elif (
553
            self.attach["fraction_increment"] is not None
554
            and self.attach["fc_final"] is not None
555
        ):
556 2
            logger.debug("Attach, Method #4")
557 2
            force_constants, targets = self._calc_method("a", self.attach, "4")
558

559 2
        elif self.attach["fc_list"] is not None:
560 2
            logger.debug("Attach, Method #5")
561 2
            force_constants, targets = self._calc_method("a", self.attach, "5")
562

563 2
        elif all(v is None for k, v in self.attach.items()):
564 2
            logger.debug("No restraint info set for the attach phase! Skipping...")
565

566
        else:
567 0
            logger.error(
568
                "Attach restraint input did not match one of the supported methods..."
569
            )
570 0
            for k, v in self.attach.items():
571 0
                logger.debug("{} = {}".format(k, v))
572 0
            raise Exception(
573
                "Attach restraint input did not match one of the supported methods..."
574
            )
575

576 2
        if force_constants is not None and targets is not None:
577 2
            self.phase["attach"]["force_constants"] = force_constants
578 2
            self.phase["attach"]["targets"] = targets
579

580
        # ------------------------------------ PULL ------------------------------------ #
581 2
        logger.debug("Calculating pull targets and force constants...")
582

583 2
        force_constants = None
584 2
        targets = None
585

586 2
        if self.auto_apr and self.pull["target_final"] is not None:
587 2
            self.pull["fc"] = self.phase["attach"]["force_constants"][-1]
588 2
            self.pull["target_initial"] = self.phase["attach"]["targets"][-1]
589

590 2
        if (
591
            self.pull["num_windows"] is not None
592
            and self.pull["target_final"] is not None
593
        ):
594 2
            if self.pull["target_initial"] is not None:
595 2
                logger.debug("Pull, Method #1")
596 2
                force_constants, targets = self._calc_method("p", self.pull, "1")
597
            else:
598 2
                logger.debug("Pull, Method #1a")
599 2
                force_constants, targets = self._calc_method("p", self.pull, "1a")
600

601 2
        elif (
602
            self.pull["target_increment"] is not None
603
            and self.pull["target_final"] is not None
604
        ):
605 2
            if self.pull["target_initial"] is not None:
606 2
                logger.debug("Pull, Method #2")
607 2
                force_constants, targets = self._calc_method("p", self.pull, "2")
608
            else:
609 2
                logger.debug("Pull, Method #2a")
610 2
                force_constants, targets = self._calc_method("p", self.pull, "2a")
611

612 2
        elif (
613
            self.pull["fraction_list"] is not None
614
            and self.pull["target_final"] is not None
615
        ):
616 2
            logger.debug("Pull, Method #3")
617 2
            force_constants, targets = self._calc_method("p", self.pull, "3")
618

619 2
        elif (
620
            self.pull["fraction_increment"] is not None
621
            and self.pull["target_final"] is not None
622
        ):
623 2
            logger.debug("Pull, Method #4")
624 2
            force_constants, targets = self._calc_method("p", self.pull, "4")
625

626 2
        elif self.pull["target_list"] is not None:
627 2
            logger.debug("Pull, Method #5")
628 2
            force_constants, targets = self._calc_method("p", self.pull, "5")
629

630 2
        elif all(v is None for k, v in self.pull.items()):
631 2
            logger.debug("No restraint info set for the pull phase! Skipping...")
632

633
        else:
634 0
            logger.error(
635
                "Pull restraint input did not match one of the supported methods..."
636
            )
637 0
            for k, v in self.pull.items():
638 0
                logger.debug("{} = {}".format(k, v))
639 0
            raise Exception(
640
                "Pull restraint input did not match one of the supported methods..."
641
            )
642

643 2
        if force_constants is not None and targets is not None:
644 2
            self.phase["pull"]["force_constants"] = force_constants
645 2
            self.phase["pull"]["targets"] = targets
646

647
        # ------------------------------------ RELEASE ------------------------------------ #
648 2
        logger.debug("Calculating release targets and force constants...")
649

650 2
        force_constants = None
651 2
        targets = None
652

653
        # I don't want auto_apr to make release restraints, unless I'm sure the user wants them.
654
        # I'm gonna assume that specifying self.attach['fc_final'] indicates you want it,
655
        # although this weakens the whole purpose of auto_apr.
656

657 2
        if self.auto_apr and self.release["fc_final"] is not None:
658 2
            self.release["target"] = self.phase["pull"]["targets"][-1]
659 2
            for key in [
660
                "fc_final",
661
                "fc_initial",
662
                "num_windows",
663
                "fc_increment",
664
                "fraction_increment",
665
                "fraction_list",
666
                "fc_list",
667
            ]:
668 2
                if self.attach[key] is not None and self.release[key] is None:
669 2
                    self.release[key] = self.attach[key]
670

671 2
        if (
672
            self.release["num_windows"] is not None
673
            and self.release["fc_final"] is not None
674
        ):
675 2
            if self.release["fc_initial"] is not None:
676 2
                logger.debug("Release, Method #1")
677 2
                force_constants, targets = self._calc_method("r", self.release, "1")
678
            else:
679 2
                logger.debug("Release, Method #1a")
680 2
                force_constants, targets = self._calc_method("r", self.release, "1a")
681

682 2
        elif (
683
            self.release["fc_increment"] is not None
684
            and self.release["fc_final"] is not None
685
        ):
686 2
            if self.release["fc_initial"] is not None:
687 2
                logger.debug("Release, Method #2")
688 2
                force_constants, targets = self._calc_method("r", self.release, "2")
689
            else:
690 2
                logger.debug("Release, Method #2a")
691 2
                force_constants, targets = self._calc_method("r", self.release, "2a")
692

693 2
        elif (
694
            self.release["fraction_list"] is not None
695
            and self.release["fc_final"] is not None
696
        ):
697 2
            logger.debug("Release, Method #3")
698 2
            force_constants, targets = self._calc_method("r", self.release, "3")
699

700 2
        elif (
701
            self.release["fraction_increment"] is not None
702
            and self.release["fc_final"] is not None
703
        ):
704 2
            logger.debug("Release, Method #4")
705 2
            force_constants, targets = self._calc_method("r", self.release, "4")
706

707 2
        elif self.release["fc_list"] is not None:
708 2
            logger.debug("Release, Method #5")
709 2
            force_constants, targets = self._calc_method("r", self.release, "5")
710

711 2
        elif all(v is None for k, v in self.release.items()):
712 2
            logger.debug("No restraint info set for the release phase! Skipping...")
713

714
        else:
715 0
            logger.error(
716
                "Release restraint input did not match one of the supported methods..."
717
            )
718 0
            for k, v in self.release.items():
719 0
                logger.debug("{} = {}".format(k, v))
720 0
            raise Exception(
721
                "Release restraint input did not match one of the supported methods..."
722
            )
723

724 2
        if force_constants is not None and targets is not None:
725 2
            self.phase["release"]["force_constants"] = force_constants
726 2
            self.phase["release"]["targets"] = targets
727

728
        # ----------------------------------- WINDOWS ------------------------------------ #
729

730 2
        for phase in ["attach", "pull", "release"]:
731 2
            if self.phase[phase]["targets"] is not None:
732 2
                window_count = len(self.phase[phase]["targets"])
733 2
                logger.debug("Number of {} windows = {}".format(phase, window_count))
734
            else:
735 2
                logger.debug(
736
                    "This restraint will be skipped in the {} phase".format(phase)
737
                )
738

739
        # ---------------------------------- ATOM MASKS ---------------------------------- #
740 2
        logger.debug("Assigning atom indices...")
741 2
        self.index1 = utils.index_from_mask(self.topology, self.mask1, self.amber_index)
742 2
        self.index2 = utils.index_from_mask(self.topology, self.mask2, self.amber_index)
743 2
        if self.mask3:
744 2
            self.index3 = utils.index_from_mask(
745
                self.topology, self.mask3, self.amber_index
746
            )
747
        else:
748 2
            self.index3 = None
749 2
        if self.mask4:
750 2
            self.index4 = utils.index_from_mask(
751
                self.topology, self.mask4, self.amber_index
752
            )
753
        else:
754 2
            self.index4 = None
755
        # If any `index` has more than one atom, mark it as a group restraint.
756 2
        if self.mask1 and len(self.index1) > 1:
757 2
            self.group1 = True
758

759 2
        if self.mask2 and len(self.index2) > 1:
760 2
            self.group2 = True
761

762 2
        if self.mask3 and len(self.index3) > 1:
763 0
            self.group3 = True
764

765 2
        if self.mask4 and len(self.index4) > 1:
766 0
            self.group4 = True
767

768

769 2
def static_DAT_restraint(
770
    restraint_mask_list,
771
    num_window_list,
772
    ref_structure,
773
    force_constant,
774
    continuous_apr=True,
775
    amber_index=False,
776
):
777
    """
778
    Create a restraint whose value does not change during a calculation.
779

780
    Parameters
781
    ----------
782
    restraint_mask_list: list
783
        A list of masks for which this restraint applies.
784
    num_window_list: list
785
        A list of windows during which this restraint will be applied, which should be in the form: [attach windows,
786
        pull windows, release windows].
787
    ref_structure: Path-like or parmed Amber object
788
        The reference structure that is used to determine the initial, **static** value for this restraint.
789
    force_constant: float
790
        The force constant for this restraint.
791
    continuous_apr: bool
792
        Whether this restraint uses ``continuous_apr``. This must be consistent with existing restraints.
793
    amber_index: bool
794
        Whether the atom indices for the restraint should be AMBER-style (+1) or not.
795

796
    Returns
797
    -------
798
    rest: ``DAT_restraint``
799
        A static restraint.
800

801
    """
802

803
    # Check num_window_list
804 0
    if len(num_window_list) != 3:
805 0
        raise ValueError(
806
            "The num_window_list needs to contain three integers corresponding to the number of windows in the "
807
            "attach, pull, and release phase, respectively "
808
        )
809

810 0
    rest = DAT_restraint()
811 0
    rest.continuous_apr = continuous_apr
812 0
    rest.amber_index = amber_index
813

814 0
    if isinstance(ref_structure, pmd.amber._amberparm.AmberParm):
815 0
        reference_trajectory = pt.load_parmed(ref_structure, traj=True)
816 0
        rest.topology = ref_structure
817 0
    elif isinstance(ref_structure, str):
818 0
        reference_trajectory = pt.iterload(ref_structure, traj=True)
819 0
        rest.topology = pmd.load_file(ref_structure, structure=True)
820
    else:
821 0
        raise TypeError(
822
            "static_DAT_restraint does not support the type associated with ref_structure:"
823
            + type(ref_structure)
824
        )
825

826 0
    rest.mask1 = restraint_mask_list[0]
827 0
    rest.mask2 = restraint_mask_list[1]
828 0
    if len(restraint_mask_list) >= 3:
829 0
        rest.mask3 = restraint_mask_list[2]
830 0
    if len(restraint_mask_list) == 4:
831 0
        rest.mask4 = restraint_mask_list[3]
832

833
    # Target value
834 0
    mask_string = " ".join(restraint_mask_list)
835 0
    if len(restraint_mask_list) == 2:
836
        # Distance restraint
837 0
        if reference_trajectory.top.has_box():
838 0
            target = pt.distance(reference_trajectory, mask_string, image=True)[0]
839 0
            logger.debug("Calculating distance with 'image = True' ...")
840
        else:
841 0
            target = pt.distance(reference_trajectory, mask_string, image=False)[0]
842 0
            logger.debug("Calculating distance with 'image = False' ...")
843 0
    elif len(restraint_mask_list) == 3:
844
        # Angle restraint
845 0
        target = pt.angle(reference_trajectory, mask_string)[0]
846 0
    elif len(restraint_mask_list) == 4:
847
        # Dihedral restraint
848 0
        target = pt.dihedral(reference_trajectory, mask_string)[0]
849
    else:
850 0
        raise IndexError(
851
            f"The number of masks -- {len(restraint_mask_list)} -- is not 2, 3, or 4 and thus is not one of the "
852
            f"supported types: distance, angle, or dihedral."
853
        )
854

855
    # Attach phase
856 0
    if num_window_list[0] is not None and num_window_list[0] != 0:
857 0
        rest.attach["target"] = target
858 0
        rest.attach["fc_initial"] = force_constant
859 0
        rest.attach["fc_final"] = force_constant
860 0
        rest.attach["num_windows"] = num_window_list[0]
861

862
    # Pull phase
863 0
    if num_window_list[1] is not None and num_window_list[1] != 0:
864 0
        rest.pull["fc"] = force_constant
865 0
        rest.pull["target_initial"] = target
866 0
        rest.pull["target_final"] = target
867 0
        rest.pull["num_windows"] = num_window_list[1]
868

869
    # Release phase
870 0
    if num_window_list[2] is not None and num_window_list[2] != 0:
871 0
        rest.release["target"] = target
872 0
        rest.release["fc_initial"] = force_constant
873 0
        rest.release["fc_final"] = force_constant
874 0
        rest.release["num_windows"] = num_window_list[2]
875

876 0
    rest.initialize()
877

878 0
    return rest
879

880

881 2
def check_restraints(restraint_list, create_window_list=False):
882
    """
883
    Do basic tests to ensure a list of ``DAT_restraint``s are consistent.
884
    This function is also overloaded to create window lists as well, which is not ideal.
885

886
    Parameters
887
    ----------
888
    restraint_list: list
889
        A list of restraints.
890
    create_window_list: bool
891
        Whether to use the restraints to create windows for the calculation.
892

893
    """
894

895 2
    if all(restraint.continuous_apr is True for restraint in restraint_list):
896 2
        logger.debug('All restraints are "continuous_apr" style.')
897 2
        all_continuous_apr = True
898 2
    elif all(restraint.continuous_apr is False for restraint in restraint_list):
899 2
        logger.debug('All restraints are not "continuous_apr" style.')
900 2
        all_continuous_apr = False
901
    else:
902 2
        raise ValueError(
903
            "All restraints must have the same setting for ``.continuous_apr``."
904
        )
905

906 2
    window_list = []
907 2
    phases = ["attach", "pull", "release"]
908 2
    for phase in phases:
909 2
        win_counts = []
910 2
        for restraint in restraint_list:
911 2
            if restraint.phase[phase]["targets"] is not None:
912 2
                win_counts.append(len(restraint.phase[phase]["targets"]))
913
            else:
914 2
                win_counts.append(0)
915 2
        max_count = np.max(win_counts)
916

917 2
        if max_count > 999:
918 0
            logger.info("Window name zero padding only applied up to 999.")
919

920
        # For each restraint, make sure the number of windows is either 0 (the restraint
921
        # is not active) or equal to the maximum number of windows for any
922
        # restraint.
923 2
        if all(count == 0 or count == max_count for count in win_counts):
924 2
            if max_count > 0:
925
                # `continuous_apr` during attach means that the final attach window
926
                # should be skipped and replaced with `p000`. `continuous_apr` during
927
                # release means that `r000` should be skipped and replaced with the
928
                # final pull window.
929

930 2
                if phase == "attach" and all_continuous_apr:
931 2
                    window_list += [
932
                        phase[0] + str("{:03.0f}".format(val))
933
                        for val in np.arange(0, max_count - 1, 1)
934
                    ]
935 2
                elif phase == "attach" and not all_continuous_apr:
936 2
                    window_list += [
937
                        phase[0] + str("{:03.0f}".format(val))
938
                        for val in np.arange(0, max_count, 1)
939
                    ]
940 2
                elif phase == "pull":
941 2
                    window_list += [
942
                        phase[0] + str("{:03.0f}".format(val))
943
                        for val in np.arange(0, max_count, 1)
944
                    ]
945 2
                elif phase == "release" and all_continuous_apr:
946 2
                    window_list += [
947
                        phase[0] + str("{:03.0f}".format(val))
948
                        for val in np.arange(1, max_count, 1)
949
                    ]
950 2
                elif phase == "release" and not all_continuous_apr:
951 2
                    window_list += [
952
                        phase[0] + str("{:03.0f}".format(val))
953
                        for val in np.arange(0, max_count, 1)
954
                    ]
955
        else:
956 2
            logger.error(
957
                "Restraints have unequal number of windows during the {} phase.".format(
958
                    phase
959
                )
960
            )
961 2
            logger.debug("Window counts for each restraint are as follows:")
962 2
            logger.debug(win_counts)
963 2
            raise Exception(
964
                "Restraints have unequal number of windows during the {} "
965
                "phase.".format(phase)
966
            )
967

968
    # Check that each restraint have at least two atoms
969 2
    for restraint in restraint_list:
970 2
        if not restraint.index1 or not restraint.index2:
971 0
            raise Exception("There must be at least two atoms in a restraint.")
972

973 2
    logger.info("Restraints appear to be consistent")
974

975 2
    if create_window_list:
976 2
        return window_list
977

978

979 2
def create_window_list(restraint_list):
980
    """
981
    Create list of APR windows. Runs everything through check_restraints because
982
    we need to do that.
983
    """
984

985 2
    return check_restraints(restraint_list, create_window_list=True)

Read our documentation on viewing source code .

Loading