Coverage for ase / constraints / hookean.py: 64.66%
116 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1import numpy as np
3from ase.constraints.constraint import FixConstraint, slice2enlist
4from ase.geometry.geometry import find_mic
7class Hookean(FixConstraint):
8 """Applies a Hookean restorative force between a pair of atoms, an atom
9 and a point, or an atom and a plane."""
11 def __init__(self, a1, a2, k, rt=None):
12 """Forces two atoms to stay close together by applying no force if
13 they are below a threshold length, rt, and applying a Hookean
14 restorative force when the distance between them exceeds rt. Can
15 also be used to tether an atom to a fixed point in space or to a
16 distance above a plane.
18 a1 : int
19 Index of atom 1
20 a2 : one of three options
21 1) index of atom 2
22 2) a fixed point in cartesian space to which to tether a1
23 3) a plane given as (A, B, C, D) in A x + B y + C z + D = 0.
24 k : float
25 Hooke's law (spring) constant to apply when distance
26 exceeds threshold_length. Units of eV A^-2.
27 rt : float
28 The threshold length below which there is no force. The
29 length is 1) between two atoms, 2) between atom and point.
30 This argument is not supplied in case 3. Units of A.
32 If a plane is specified, the Hooke's law force is applied if the atom
33 is on the normal side of the plane. For instance, the plane with
34 (A, B, C, D) = (0, 0, 1, -7) defines a plane in the xy plane with a z
35 intercept of +7 and a normal vector pointing in the +z direction.
36 If the atom has z > 7, then a downward force would be applied of
37 k * (atom.z - 7). The same plane with the normal vector pointing in
38 the -z direction would be given by (A, B, C, D) = (0, 0, -1, 7).
40 References
41 ----------
43 Andrew A. Peterson, Topics in Catalysis volume 57, pages40–53 (2014)
44 https://link.springer.com/article/10.1007%2Fs11244-013-0161-8
45 """
47 if isinstance(a2, int):
48 self._type = 'two atoms'
49 self.indices = [a1, a2]
50 elif len(a2) == 3:
51 self._type = 'point'
52 self.index = a1
53 self.origin = np.array(a2)
54 elif len(a2) == 4:
55 self._type = 'plane'
56 self.index = a1
57 self.plane = a2
58 else:
59 raise RuntimeError('Unknown type for a2')
60 self.threshold = rt
61 self.spring = k
63 def get_removed_dof(self, atoms):
64 return 0
66 def todict(self):
67 dct = {'name': 'Hookean'}
68 dct['kwargs'] = {'rt': self.threshold, 'k': self.spring}
69 if self._type == 'two atoms':
70 dct['kwargs']['a1'] = self.indices[0]
71 dct['kwargs']['a2'] = self.indices[1]
72 elif self._type == 'point':
73 dct['kwargs']['a1'] = self.index
74 dct['kwargs']['a2'] = self.origin
75 elif self._type == 'plane':
76 dct['kwargs']['a1'] = self.index
77 dct['kwargs']['a2'] = self.plane
78 else:
79 raise NotImplementedError(f'Bad type: {self._type}')
80 return dct
82 def adjust_positions(self, atoms, newpositions):
83 pass
85 def adjust_momenta(self, atoms, momenta):
86 pass
88 def adjust_forces(self, atoms, forces):
89 positions = atoms.positions
90 if self._type == 'plane':
91 A, B, C, D = self.plane
92 x, y, z = positions[self.index]
93 d = (A * x + B * y + C * z + D) / np.sqrt(A**2 + B**2 + C**2)
94 if d < 0:
95 return
96 magnitude = self.spring * d
97 direction = -np.array((A, B, C)) / np.linalg.norm((A, B, C))
98 forces[self.index] += direction * magnitude
99 return
100 if self._type == 'two atoms':
101 p1, p2 = positions[self.indices]
102 elif self._type == 'point':
103 p1 = positions[self.index]
104 p2 = self.origin
105 displace, _ = find_mic(p2 - p1, atoms.cell, atoms.pbc)
106 bondlength = np.linalg.norm(displace)
107 if bondlength > self.threshold:
108 magnitude = self.spring * (bondlength - self.threshold)
109 direction = displace / np.linalg.norm(displace)
110 if self._type == 'two atoms':
111 forces[self.indices[0]] += direction * magnitude
112 forces[self.indices[1]] -= direction * magnitude
113 else:
114 forces[self.index] += direction * magnitude
116 def adjust_potential_energy(self, atoms):
117 """Returns the difference to the potential energy due to an active
118 constraint. (That is, the quantity returned is to be added to the
119 potential energy.)"""
120 positions = atoms.positions
121 if self._type == 'plane':
122 A, B, C, D = self.plane
123 x, y, z = positions[self.index]
124 d = (A * x + B * y + C * z + D) / np.sqrt(A**2 + B**2 + C**2)
125 if d > 0:
126 return 0.5 * self.spring * d**2
127 else:
128 return 0.0
129 if self._type == 'two atoms':
130 p1, p2 = positions[self.indices]
131 elif self._type == 'point':
132 p1 = positions[self.index]
133 p2 = self.origin
134 displace, _ = find_mic(p2 - p1, atoms.cell, atoms.pbc)
135 bondlength = np.linalg.norm(displace)
136 if bondlength > self.threshold:
137 return 0.5 * self.spring * (bondlength - self.threshold) ** 2
138 else:
139 return 0.0
141 def get_indices(self):
142 if self._type == 'two atoms':
143 return self.indices
144 elif self._type == 'point':
145 return self.index
146 elif self._type == 'plane':
147 return self.index
149 def index_shuffle(self, atoms, ind):
150 # See docstring of superclass
151 if self._type == 'two atoms':
152 newa = [-1, -1] # Signal error
153 for new, old in slice2enlist(ind, len(atoms)):
154 for i, a in enumerate(self.indices):
155 if old == a:
156 newa[i] = new
157 if newa[0] == -1 or newa[1] == -1:
158 raise IndexError('Constraint not part of slice')
159 self.indices = newa
160 elif (self._type == 'point') or (self._type == 'plane'):
161 newa = -1 # Signal error
162 for new, old in slice2enlist(ind, len(atoms)):
163 if old == self.index:
164 newa = new
165 break
166 if newa == -1:
167 raise IndexError('Constraint not part of slice')
168 self.index = newa
170 def __repr__(self):
171 if self._type == 'two atoms':
172 return 'Hookean(%d, %d)' % tuple(self.indices)
173 elif self._type == 'point':
174 return 'Hookean(%d) to cartesian' % self.index
175 else:
176 return 'Hookean(%d) to plane' % self.index