Coverage for /builds/ase/ase/ase/calculators/combine_mm.py: 87.92%

1# fmt: off 

2 

3import copy 

4 

5import numpy as np 

6 

7from ase import units 

8from ase.calculators.calculator import Calculator 

9from ase.calculators.qmmm import combine_lj_lorenz_berthelot 

10 

11k_c = units.Hartree * units.Bohr 

12 

13 

14class CombineMM(Calculator): 

15 implemented_properties = ['energy', 'forces'] 

16 

17 def __init__(self, idx, apm1, apm2, calc1, calc2, 

18 sig1, eps1, sig2, eps2, rc=7.0, width=1.0): 

19 """A calculator that combines two MM calculators 

20 (TIPnP, Counterions, ...) 

21 

22 parameters: 

23 

24 idx: List of indices of atoms belonging to calculator 1 

25 apm1,2: atoms pr molecule of each subsystem (NB: apm for TIP4P is 3!) 

26 calc1,2: calculator objects for each subsystem 

27 sig1,2, eps1,2: LJ parameters for each subsystem. Should be a numpy 

28 array of length = apm 

29 rc = long range cutoff 

30 width = width of cutoff region. 

31 

32 Currently the interactions are limited to being: 

33 - Nonbonded 

34 - Hardcoded to two terms: 

35 - Coulomb electrostatics 

36 - Lennard-Jones 

37 

38 It could of course benefit from being more like the EIQMMM class 

39 where the interactions are switchable. But this is in princple 

40 just meant for adding counter ions to a qmmm simulation to neutralize 

41 the charge of the total systemn 

42 

43 Maybe it can combine n MM calculators in the future? 

44 """ 

45 

46 self.idx = idx 

47 self.apm1 = apm1 # atoms per mol for LJ calculator 

48 self.apm2 = apm2 

49 

50 self.rc = rc 

51 self.width = width 

52 

53 self.atoms1 = None 

54 self.atoms2 = None 

55 self.mask = None 

56 

57 self.calc1 = calc1 

58 self.calc2 = calc2 

59 

60 self.sig1 = sig1 

61 self.eps1 = eps1 

62 self.sig2 = sig2 

63 self.eps2 = eps2 

64 

65 Calculator.__init__(self) 

66 

67 def initialize(self, atoms): 

68 self.mask = np.zeros(len(atoms), bool) 

69 self.mask[self.idx] = True 

70 

71 constraints = atoms.constraints 

72 atoms.constraints = [] 

73 self.atoms1 = atoms[self.mask] 

74 self.atoms2 = atoms[~self.mask] 

75 

76 atoms.constraints = constraints 

77 

78 self.atoms1.calc = self.calc1 

79 self.atoms2.calc = self.calc2 

80 

81 self.cell = atoms.cell 

82 self.pbc = atoms.pbc 

83 

84 self.sigma, self.epsilon =\ 

85 combine_lj_lorenz_berthelot(self.sig1, self.sig2, 

86 self.eps1, self.eps2) 

87 

88 self.make_virtual_mask() 

89 

90 def calculate(self, atoms, properties, system_changes): 

91 Calculator.calculate(self, atoms, properties, system_changes) 

92 

93 if self.atoms1 is None: 

94 self.initialize(atoms) 

95 

96 pos1 = atoms.positions[self.mask] 

97 pos2 = atoms.positions[~self.mask] 

98 self.atoms1.set_positions(pos1) 

99 self.atoms2.set_positions(pos2) 

100 

101 # positions and charges for the coupling term, which should 

102 # include virtual charges and sites: 

103 spm1 = self.atoms1.calc.sites_per_mol 

104 spm2 = self.atoms2.calc.sites_per_mol 

105 xpos1 = self.atoms1.calc.add_virtual_sites(pos1) 

106 xpos2 = self.atoms2.calc.add_virtual_sites(pos2) 

107 

108 xc1 = self.atoms1.calc.get_virtual_charges(self.atoms1) 

109 xc2 = self.atoms2.calc.get_virtual_charges(self.atoms2) 

110 

111 xpos1 = xpos1.reshape((-1, spm1, 3)) 

112 xpos2 = xpos2.reshape((-1, spm2, 3)) 

113 

114 e_c, f_c = self.coulomb(xpos1, xpos2, xc1, xc2, spm1, spm2) 

115 

116 e_lj, f1, f2 = self.lennard_jones(self.atoms1, self.atoms2) 

117 

118 f_lj = np.zeros((len(atoms), 3)) 

119 f_lj[self.mask] += f1 

120 f_lj[~self.mask] += f2 

121 

122 # internal energy, forces of each subsystem: 

123 f12 = np.zeros((len(atoms), 3)) 

124 e1 = self.atoms1.get_potential_energy() 

125 fi1 = self.atoms1.get_forces() 

126 

127 e2 = self.atoms2.get_potential_energy() 

128 fi2 = self.atoms2.get_forces() 

129 

130 f12[self.mask] += fi1 

131 f12[~self.mask] += fi2 

132 

133 self.results['energy'] = e_c + e_lj + e1 + e2 

134 self.results['forces'] = f_c + f_lj + f12 

135 

136 def get_virtual_charges(self, atoms): 

137 if self.atoms1 is None: 

138 self.initialize(atoms) 

139 

140 vc1 = self.atoms1.calc.get_virtual_charges(atoms[self.mask]) 

141 vc2 = self.atoms2.calc.get_virtual_charges(atoms[~self.mask]) 

142 # Need to expand mask with possible new virtual sites. 

143 # Virtual sites should ALWAYS be put AFTER actual atoms, like in 

144 # TIP4P: OHHX, OHHX, ... 

145 

146 vc = np.zeros(len(vc1) + len(vc2)) 

147 vc[self.virtual_mask] = vc1 

148 vc[~self.virtual_mask] = vc2 

149 

150 return vc 

151 

152 def add_virtual_sites(self, positions): 

153 vs1 = self.atoms1.calc.add_virtual_sites(positions[self.mask]) 

154 vs2 = self.atoms2.calc.add_virtual_sites(positions[~self.mask]) 

155 vs = np.zeros((len(vs1) + len(vs2), 3)) 

156 

157 vs[self.virtual_mask] = vs1 

158 vs[~self.virtual_mask] = vs2 

159 

160 return vs 

161 

162 def make_virtual_mask(self): 

163 virtual_mask = [] 

164 ct1 = 0 

165 ct2 = 0 

166 for i in range(len(self.mask)): 

167 virtual_mask.append(self.mask[i]) 

168 if self.mask[i]: 

169 ct1 += 1 

170 if not self.mask[i]: 

171 ct2 += 1 

172 if ((ct2 == self.apm2) & 

173 (self.apm2 != self.atoms2.calc.sites_per_mol)): 

174 virtual_mask.append(False) 

175 ct2 = 0 

176 if ((ct1 == self.apm1) & 

177 (self.apm1 != self.atoms1.calc.sites_per_mol)): 

178 virtual_mask.append(True) 

179 ct1 = 0 

180 

181 self.virtual_mask = np.array(virtual_mask) 

182 

183 def coulomb(self, xpos1, xpos2, xc1, xc2, spm1, spm2): 

184 energy = 0.0 

185 forces = np.zeros((len(xc1) + len(xc2), 3)) 

186 

187 self.xpos1 = xpos1 

188 self.xpos2 = xpos2 

189 

190 R1 = xpos1 

191 R2 = xpos2 

192 F1 = np.zeros_like(R1) 

193 F2 = np.zeros_like(R2) 

194 C1 = xc1.reshape((-1, np.shape(xpos1)[1])) 

195 C2 = xc2.reshape((-1, np.shape(xpos2)[1])) 

196 # Vectorized evaluation is not as trivial when spm1 != spm2. 

197 # This is pretty inefficient, but for ~1-5 counter ions as region 1 

198 # it should not matter much .. 

199 # There is definitely room for improvements here. 

200 cell = self.cell.diagonal() 

201 for m1, (r1, c1) in enumerate(zip(R1, C1)): 

202 for m2, (r2, c2) in enumerate(zip(R2, C2)): 

203 r00 = r2[0] - r1[0] 

204 shift = np.zeros(3) 

205 for i, periodic in enumerate(self.pbc): 

206 if periodic: 

207 L = cell[i] 

208 shift[i] = (r00[i] + L / 2.) % L - L / 2. - r00[i] 

209 r00 += shift 

210 

211 d00 = (r00**2).sum()**0.5 

212 t = 1 

213 dtdd = 0 

214 if d00 > self.rc: 

215 continue 

216 elif d00 > self.rc - self.width: 

217 y = (d00 - self.rc + self.width) / self.width 

218 t -= y**2 * (3.0 - 2.0 * y) 

219 dtdd = r00 * 6 * y * (1.0 - y) / (self.width * d00) 

220 

221 for a1 in range(spm1): 

222 for a2 in range(spm2): 

223 r = r2[a2] - r1[a1] + shift 

224 d2 = (r**2).sum() 

225 d = d2**0.5 

226 e = k_c * c1[a1] * c2[a2] / d 

227 energy += t * e 

228 

229 F1[m1, a1] -= t * (e / d2) * r 

230 F2[m2, a2] += t * (e / d2) * r 

231 

232 F1[m1, 0] -= dtdd * e 

233 F2[m2, 0] += dtdd * e 

234 

235 F1 = F1.reshape((-1, 3)) 

236 F2 = F2.reshape((-1, 3)) 

237 

238 # Redist forces but dont save forces in org calculators 

239 atoms1 = self.atoms1.copy() 

240 atoms1.calc = copy.copy(self.calc1) 

241 atoms1.calc.atoms = atoms1 

242 F1 = atoms1.calc.redistribute_forces(F1) 

243 atoms2 = self.atoms2.copy() 

244 atoms2.calc = copy.copy(self.calc2) 

245 atoms2.calc.atoms = atoms2 

246 F2 = atoms2.calc.redistribute_forces(F2) 

247 

248 forces = np.zeros((len(self.atoms), 3)) 

249 forces[self.mask] = F1 

250 forces[~self.mask] = F2 

251 return energy, forces 

252 

253 def lennard_jones(self, atoms1, atoms2): 

254 pos1 = atoms1.get_positions().reshape((-1, self.apm1, 3)) 

255 pos2 = atoms2.get_positions().reshape((-1, self.apm2, 3)) 

256 

257 f1 = np.zeros_like(atoms1.positions) 

258 f2 = np.zeros_like(atoms2.positions) 

259 energy = 0.0 

260 

261 cell = self.cell.diagonal() 

262 for q, p1 in enumerate(pos1): # molwise loop 

263 eps = self.epsilon 

264 sig = self.sigma 

265 

266 R00 = pos2[:, 0] - p1[0, :] 

267 

268 # cutoff from first atom of each mol 

269 shift = np.zeros_like(R00) 

270 for i, periodic in enumerate(self.pbc): 

271 if periodic: 

272 L = cell[i] 

273 shift[:, i] = (R00[:, i] + L / 2) % L - L / 2 - R00[:, i] 

274 R00 += shift 

275 

276 d002 = (R00**2).sum(1) 

277 d00 = d002**0.5 

278 x1 = d00 > self.rc - self.width 

279 x2 = d00 < self.rc 

280 x12 = np.logical_and(x1, x2) 

281 y = (d00[x12] - self.rc + self.width) / self.width 

282 t = np.zeros(len(d00)) 

283 t[x2] = 1.0 

284 t[x12] -= y**2 * (3.0 - 2.0 * y) 

285 dt = np.zeros(len(d00)) 

286 dt[x12] -= 6.0 / self.width * y * (1.0 - y) 

287 for qa in range(len(p1)): 

288 if ~np.any(eps[qa, :]): 

289 continue 

290 R = pos2 - p1[qa, :] + shift[:, None] 

291 d2 = (R**2).sum(2) 

292 c6 = (sig[qa, :]**2 / d2)**3 

293 c12 = c6**2 

294 e = 4 * eps[qa, :] * (c12 - c6) 

295 energy += np.dot(e.sum(1), t) 

296 f = t[:, None, None] * (24 * eps[qa, :] * 

297 (2 * c12 - c6) / d2)[:, :, None] * R 

298 f00 = - (e.sum(1) * dt / d00)[:, None] * R00 

299 f2 += f.reshape((-1, 3)) 

300 f1[q * self.apm1 + qa, :] -= f.sum(0).sum(0) 

301 f1[q * self.apm1, :] -= f00.sum(0) 

302 f2[::self.apm2, :] += f00 

303 

304 return energy, f1, f2 

305 

306 def redistribute_forces(self, forces): 

307 f1 = self.calc1.redistribute_forces(forces[self.virtual_mask]) 

308 f2 = self.calc2.redistribute_forces(forces[~self.virtual_mask]) 

309 # and then they are back on the real atom centers so 

310 f = np.zeros((len(self.atoms), 3)) 

311 f[self.mask] = f1 

312 f[~self.mask] = f2 

313 return f