HarmonicAngleJaxForce

Source code in dmff/classical/intra.py
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
class HarmonicAngleJaxForce:
    def __init__(self, p1idx, p2idx, p3idx, prmidx):
        self.p1idx = p1idx
        self.p2idx = p2idx
        self.p3idx = p3idx
        self.prmidx = prmidx
        self.refresh_calculators()

    def generate_get_energy(self):
        def get_energy(positions, box, pairs, k, theta0):
            p1 = positions[self.p1idx,:]
            p2 = positions[self.p2idx,:]
            p3 = positions[self.p3idx,:]
            kprm = k[self.prmidx]
            theta0prm = theta0[self.prmidx]
            ang = angle(p1, p2, p3)
            return jnp.sum(0.5 * kprm * jnp.power(ang - theta0prm, 2))

        return get_energy

    def update_env(self, attr, val):
        """
        Update the environment of the calculator
        """
        setattr(self, attr, val)
        self.refresh_calculators()

    def refresh_calculators(self):
        """
        refresh the energy and force calculators according to the current environment
        """
        self.get_energy = self.generate_get_energy()
        self.get_forces = value_and_grad(self.get_energy)

refresh_calculators()

refresh the energy and force calculators according to the current environment

Source code in dmff/classical/intra.py
92
93
94
95
96
97
def refresh_calculators(self):
    """
    refresh the energy and force calculators according to the current environment
    """
    self.get_energy = self.generate_get_energy()
    self.get_forces = value_and_grad(self.get_energy)

update_env(attr, val)

Update the environment of the calculator

Source code in dmff/classical/intra.py
85
86
87
88
89
90
def update_env(self, attr, val):
    """
    Update the environment of the calculator
    """
    setattr(self, attr, val)
    self.refresh_calculators()

HarmonicBondJaxForce

Source code in dmff/classical/intra.py
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class HarmonicBondJaxForce:
    def __init__(self, p1idx, p2idx, prmidx):
        self.p1idx = p1idx
        self.p2idx = p2idx
        self.prmidx = prmidx
        self.refresh_calculators()

    def generate_get_energy(self):
        def get_energy(positions, box, pairs, k, length):
            p1 = positions[self.p1idx,:]
            p2 = positions[self.p2idx,:]
            kprm = k[self.prmidx]
            b0prm = length[self.prmidx]
            dist = distance(p1, p2)
            return jnp.sum(0.5 * kprm * jnp.power(dist - b0prm, 2))

        return get_energy

    def update_env(self, attr, val):
        """
        Update the environment of the calculator
        """
        setattr(self, attr, val)
        self.refresh_calculators()

    def refresh_calculators(self):
        """
        refresh the energy and force calculators according to the current environment
        """
        self.get_energy = self.generate_get_energy()
        self.get_forces = value_and_grad(self.get_energy)

refresh_calculators()

refresh the energy and force calculators according to the current environment

Source code in dmff/classical/intra.py
57
58
59
60
61
62
def refresh_calculators(self):
    """
    refresh the energy and force calculators according to the current environment
    """
    self.get_energy = self.generate_get_energy()
    self.get_forces = value_and_grad(self.get_energy)

update_env(attr, val)

Update the environment of the calculator

Source code in dmff/classical/intra.py
50
51
52
53
54
55
def update_env(self, attr, val):
    """
    Update the environment of the calculator
    """
    setattr(self, attr, val)
    self.refresh_calculators()

PeriodicTorsionJaxForce

Source code in dmff/classical/intra.py
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
class PeriodicTorsionJaxForce:
    def __init__(self, p1idx, p2idx, p3idx, p4idx, prmidx, order):
        self.p1idx = p1idx
        self.p2idx = p2idx
        self.p3idx = p3idx
        self.p4idx = p4idx
        self.prmidx = prmidx
        self.order = order
        self.refresh_calculators()

    def generate_get_energy(self):
        def get_energy(positions, box, pairs, k, psi):
            p1 = positions[self.p1idx,:]
            p2 = positions[self.p2idx,:]
            p3 = positions[self.p3idx,:]
            p4 = positions[self.p4idx,:]
            kp = k[self.prmidx]
            psip = psi[self.prmidx]
            dih = dihedral(p1, p2, p3, p4)
            ener = kp * (1 + jnp.cos(self.order * dih - psip))
            return jnp.sum(ener)

        return get_energy

    def update_env(self, attr, val):
        """
        Update the environment of the calculator
        """
        setattr(self, attr, val)
        self.refresh_calculators()

    def refresh_calculators(self):
        """
        refresh the energy and force calculators according to the current environment
        """
        self.get_energy = self.generate_get_energy()
        self.get_forces = value_and_grad(self.get_energy)

refresh_calculators()

refresh the energy and force calculators according to the current environment

Source code in dmff/classical/intra.py
131
132
133
134
135
136
def refresh_calculators(self):
    """
    refresh the energy and force calculators according to the current environment
    """
    self.get_energy = self.generate_get_energy()
    self.get_forces = value_and_grad(self.get_energy)

update_env(attr, val)

Update the environment of the calculator

Source code in dmff/classical/intra.py
124
125
126
127
128
129
def update_env(self, attr, val):
    """
    Update the environment of the calculator
    """
    setattr(self, attr, val)
    self.refresh_calculators()