Coverage for ase / constraints / fix_linear_triatomic.py: 97.47%

1import numpy as np 

2 

3from ase.constraints.constraint import FixConstraint 

4from ase.geometry import find_mic, wrap_positions 

5 

6 

7class FixLinearTriatomic(FixConstraint): 

8 """Holonomic constraints for rigid linear triatomic molecules.""" 

9 

10 def __init__(self, triples): 

11 """Apply RATTLE-type bond constraints between outer atoms n and m 

12 and linear vectorial constraints to the position of central 

13 atoms o to fix the geometry of linear triatomic molecules of the 

14 type: 

15 

16 n--o--m 

17 

18 Parameters 

19 ---------- 

20 

21 triples: list 

22 Indices of the atoms forming the linear molecules to constrain 

23 as triples. Sequence should be (n, o, m) or (m, o, n). 

24 

25 When using these constraints in molecular dynamics or structure 

26 optimizations, atomic forces need to be redistributed within a 

27 triple. The function redistribute_forces_optimization implements 

28 the redistribution of forces for structure optimization, while 

29 the function redistribute_forces_md implements the redistribution 

30 for molecular dynamics. 

31 

32 References 

33 ---------- 

34 

35 Ciccotti et al. Molecular Physics 47 (1982) 

36 :doi:`10.1080/00268978200100942` 

37 """ 

38 self.triples = np.asarray(triples) 

39 if self.triples.shape[1] != 3: 

40 raise ValueError('"triples" has wrong size') 

41 self.bondlengths = None 

42 

43 def get_removed_dof(self, atoms): 

44 return 4 * len(self.triples) 

45 

46 @property 

47 def n_ind(self): 

48 return self.triples[:, 0] 

49 

50 @property 

51 def m_ind(self): 

52 return self.triples[:, 2] 

53 

54 @property 

55 def o_ind(self): 

56 return self.triples[:, 1] 

57 

58 def initialize(self, atoms): 

59 masses = atoms.get_masses() 

60 self.mass_n, self.mass_m, self.mass_o = self.get_slices(masses) 

61 

62 self.bondlengths = self.initialize_bond_lengths(atoms) 

63 self.bondlengths_nm = self.bondlengths.sum(axis=1) 

64 

65 C1 = self.bondlengths[:, ::-1] / self.bondlengths_nm[:, None] 

66 C2 = ( 

67 C1[:, 0] ** 2 * self.mass_o * self.mass_m 

68 + C1[:, 1] ** 2 * self.mass_n * self.mass_o 

69 + self.mass_n * self.mass_m 

70 ) 

71 C2 = C1 / C2[:, None] 

72 C3 = self.mass_n * C1[:, 1] - self.mass_m * C1[:, 0] 

73 C3 = C2 * self.mass_o[:, None] * C3[:, None] 

74 C3[:, 1] *= -1 

75 C3 = (C3 + 1) / np.vstack((self.mass_n, self.mass_m)).T 

76 C4 = C1[:, 0] ** 2 + C1[:, 1] ** 2 + 1 

77 C4 = C1 / C4[:, None] 

78 

79 self.C1 = C1 

80 self.C2 = C2 

81 self.C3 = C3 

82 self.C4 = C4 

83 

84 def adjust_positions(self, atoms, new): 

85 old = atoms.positions 

86 new_n, new_m, new_o = self.get_slices(new) 

87 

88 if self.bondlengths is None: 

89 self.initialize(atoms) 

90 

91 r0 = old[self.n_ind] - old[self.m_ind] 

92 d0, _ = find_mic(r0, atoms.cell, atoms.pbc) 

93 d1 = new_n - new_m - r0 + d0 

94 a = np.einsum('ij,ij->i', d0, d0) 

95 b = np.einsum('ij,ij->i', d1, d0) 

96 c = np.einsum('ij,ij->i', d1, d1) - self.bondlengths_nm**2 

97 g = (b - (b**2 - a * c) ** 0.5) / (a * self.C3.sum(axis=1)) 

98 g = g[:, None] * self.C3 

99 new_n -= g[:, 0, None] * d0 

100 new_m += g[:, 1, None] * d0 

101 if np.allclose(d0, r0): 

102 new_o = self.C1[:, 0, None] * new_n + self.C1[:, 1, None] * new_m 

103 else: 

104 v1, _ = find_mic(new_n, atoms.cell, atoms.pbc) 

105 v2, _ = find_mic(new_m, atoms.cell, atoms.pbc) 

106 rb = self.C1[:, 0, None] * v1 + self.C1[:, 1, None] * v2 

107 new_o = wrap_positions(rb, atoms.cell, atoms.pbc) 

108 

109 self.set_slices(new_n, new_m, new_o, new) 

110 

111 def adjust_momenta(self, atoms, p): 

112 old = atoms.positions 

113 p_n, p_m, p_o = self.get_slices(p) 

114 

115 if self.bondlengths is None: 

116 self.initialize(atoms) 

117 

118 mass_nn = self.mass_n[:, None] 

119 mass_mm = self.mass_m[:, None] 

120 mass_oo = self.mass_o[:, None] 

121 

122 d = old[self.n_ind] - old[self.m_ind] 

123 d, _ = find_mic(d, atoms.cell, atoms.pbc) 

124 dv = p_n / mass_nn - p_m / mass_mm 

125 k = np.einsum('ij,ij->i', dv, d) / self.bondlengths_nm**2 

126 k = self.C3 / (self.C3.sum(axis=1)[:, None]) * k[:, None] 

127 p_n -= k[:, 0, None] * mass_nn * d 

128 p_m += k[:, 1, None] * mass_mm * d 

129 p_o = mass_oo * ( 

130 self.C1[:, 0, None] * p_n / mass_nn 

131 + self.C1[:, 1, None] * p_m / mass_mm 

132 ) 

133 

134 self.set_slices(p_n, p_m, p_o, p) 

135 

136 def adjust_forces(self, atoms, forces): 

137 if self.bondlengths is None: 

138 self.initialize(atoms) 

139 

140 A = self.C4 * np.diff(self.C1) 

141 A[:, 0] *= -1 

142 A -= 1 

143 B = np.diff(self.C4) / (A.sum(axis=1))[:, None] 

144 A /= (A.sum(axis=1))[:, None] 

145 

146 self.constraint_forces = -forces 

147 old = atoms.positions 

148 

149 fr_n, fr_m, fr_o = self.redistribute_forces_optimization(forces) 

150 

151 d = old[self.n_ind] - old[self.m_ind] 

152 d, _ = find_mic(d, atoms.cell, atoms.pbc) 

153 df = fr_n - fr_m 

154 k = -np.einsum('ij,ij->i', df, d) / self.bondlengths_nm**2 

155 forces[self.n_ind] = fr_n + k[:, None] * d * A[:, 0, None] 

156 forces[self.m_ind] = fr_m - k[:, None] * d * A[:, 1, None] 

157 forces[self.o_ind] = fr_o + k[:, None] * d * B 

158 

159 self.constraint_forces += forces 

160 

161 def redistribute_forces_optimization(self, forces): 

162 """Redistribute forces within a triple when performing structure 

163 optimizations. 

164 

165 The redistributed forces needs to be further adjusted using the 

166 appropriate Lagrange multipliers as implemented in adjust_forces.""" 

167 forces_n, forces_m, forces_o = self.get_slices(forces) 

168 C1_1 = self.C1[:, 0, None] 

169 C1_2 = self.C1[:, 1, None] 

170 C4_1 = self.C4[:, 0, None] 

171 C4_2 = self.C4[:, 1, None] 

172 

173 fr_n = (1 - C4_1 * C1_1) * forces_n - C4_1 * ( 

174 C1_2 * forces_m - forces_o 

175 ) 

176 fr_m = (1 - C4_2 * C1_2) * forces_m - C4_2 * ( 

177 C1_1 * forces_n - forces_o 

178 ) 

179 fr_o = ( 

180 (1 - 1 / (C1_1**2 + C1_2**2 + 1)) * forces_o 

181 + C4_1 * forces_n 

182 + C4_2 * forces_m 

183 ) 

184 

185 return fr_n, fr_m, fr_o 

186 

187 def redistribute_forces_md(self, atoms, forces, rand=False): 

188 """Redistribute forces within a triple when performing molecular 

189 dynamics. 

190 

191 When rand=True, use the equations for random force terms, as 

192 used e.g. by Langevin dynamics, otherwise apply the standard 

193 equations for deterministic forces (see Ciccotti et al. Molecular 

194 Physics 47 (1982)).""" 

195 if self.bondlengths is None: 

196 self.initialize(atoms) 

197 forces_n, forces_m, forces_o = self.get_slices(forces) 

198 C1_1 = self.C1[:, 0, None] 

199 C1_2 = self.C1[:, 1, None] 

200 C2_1 = self.C2[:, 0, None] 

201 C2_2 = self.C2[:, 1, None] 

202 mass_nn = self.mass_n[:, None] 

203 mass_mm = self.mass_m[:, None] 

204 mass_oo = self.mass_o[:, None] 

205 if rand: 

206 mr1 = (mass_mm / mass_nn) ** 0.5 

207 mr2 = (mass_oo / mass_nn) ** 0.5 

208 mr3 = (mass_nn / mass_mm) ** 0.5 

209 mr4 = (mass_oo / mass_mm) ** 0.5 

210 else: 

211 mr1 = 1.0 

212 mr2 = 1.0 

213 mr3 = 1.0 

214 mr4 = 1.0 

215 

216 fr_n = (1 - C1_1 * C2_1 * mass_oo * mass_mm) * forces_n - C2_1 * ( 

217 C1_2 * mr1 * mass_oo * mass_nn * forces_m 

218 - mr2 * mass_mm * mass_nn * forces_o 

219 ) 

220 

221 fr_m = (1 - C1_2 * C2_2 * mass_oo * mass_nn) * forces_m - C2_2 * ( 

222 C1_1 * mr3 * mass_oo * mass_mm * forces_n 

223 - mr4 * mass_mm * mass_nn * forces_o 

224 ) 

225 

226 self.set_slices(fr_n, fr_m, 0.0, forces) 

227 

228 def get_slices(self, a): 

229 a_n = a[self.n_ind] 

230 a_m = a[self.m_ind] 

231 a_o = a[self.o_ind] 

232 

233 return a_n, a_m, a_o 

234 

235 def set_slices(self, a_n, a_m, a_o, a): 

236 a[self.n_ind] = a_n 

237 a[self.m_ind] = a_m 

238 a[self.o_ind] = a_o 

239 

240 def initialize_bond_lengths(self, atoms): 

241 bondlengths = np.zeros((len(self.triples), 2)) 

242 

243 for i in range(len(self.triples)): 

244 bondlengths[i, 0] = atoms.get_distance( 

245 self.n_ind[i], self.o_ind[i], mic=True 

246 ) 

247 bondlengths[i, 1] = atoms.get_distance( 

248 self.o_ind[i], self.m_ind[i], mic=True 

249 ) 

250 

251 return bondlengths 

252 

253 def get_indices(self): 

254 return np.unique(self.triples.ravel()) 

255 

256 def todict(self): 

257 return { 

258 'name': 'FixLinearTriatomic', 

259 'kwargs': {'triples': self.triples.tolist()}, 

260 } 

261 

262 def index_shuffle(self, atoms, ind): 

263 """Shuffle the indices of the three atoms in this constraint""" 

264 map = np.zeros(len(atoms), int) 

265 map[ind] = 1 

266 n = map.sum() 

267 map[:] = -1 

268 map[ind] = range(n) 

269 triples = map[self.triples] 

270 self.triples = triples[(triples != -1).all(1)] 

271 if len(self.triples) == 0: 

272 raise IndexError('Constraint not part of slice')