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

1import copy 

2 

3import numpy as np 

4 

5from ase import units 

6from ase.calculators.calculator import Calculator 

7from ase.calculators.qmmm import combine_lj_lorenz_berthelot 

8 

9k_c = units.Hartree * units.Bohr 

10 

11 

12class CombineMM(Calculator): 

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

14 

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

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

17 """A calculator that combines two MM calculators 

18 (TIPnP, Counterions, ...) 

19 

20 parameters: 

21 

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

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

24 calc1,2: calculator objects for each subsystem 

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

26 array of length = apm 

27 rc = long range cutoff 

28 width = width of cutoff region. 

29 

30 Currently the interactions are limited to being: 

31 - Nonbonded 

32 - Hardcoded to two terms: 

33 - Coulomb electrostatics 

34 - Lennard-Jones 

35 

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

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

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

39 the charge of the total systemn 

40 

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

42 """ 

43 

44 self.idx = idx 

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

46 self.apm2 = apm2 

47 

48 self.rc = rc 

49 self.width = width 

50 

51 self.atoms1 = None 

52 self.atoms2 = None 

53 self.mask = None 

54 

55 self.calc1 = calc1 

56 self.calc2 = calc2 

57 

58 self.sig1 = sig1 

59 self.eps1 = eps1 

60 self.sig2 = sig2 

61 self.eps2 = eps2 

62 

63 Calculator.__init__(self) 

64 

65 def initialize(self, atoms): 

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

67 self.mask[self.idx] = True 

68 

69 constraints = atoms.constraints 

70 atoms.constraints = [] 

71 self.atoms1 = atoms[self.mask] 

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

73 

74 atoms.constraints = constraints 

75 

76 self.atoms1.calc = self.calc1 

77 self.atoms2.calc = self.calc2 

78 

79 self.cell = atoms.cell 

80 self.pbc = atoms.pbc 

81 

82 self.sigma, self.epsilon =\ 

83 combine_lj_lorenz_berthelot(self.sig1, self.sig2, 

84 self.eps1, self.eps2) 

85 

86 self.make_virtual_mask() 

87 

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

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

90 

91 if self.atoms1 is None: 

92 self.initialize(atoms) 

93 

94 pos1 = atoms.positions[self.mask] 

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

96 self.atoms1.set_positions(pos1) 

97 self.atoms2.set_positions(pos2) 

98 

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

100 # include virtual charges and sites: 

101 spm1 = self.atoms1.calc.sites_per_mol 

102 spm2 = self.atoms2.calc.sites_per_mol 

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

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

105 

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

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

108 

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

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

111 

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

113 

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

115 

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

117 f_lj[self.mask] += f1 

118 f_lj[~self.mask] += f2 

119 

120 # internal energy, forces of each subsystem: 

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

122 e1 = self.atoms1.get_potential_energy() 

123 fi1 = self.atoms1.get_forces() 

124 

125 e2 = self.atoms2.get_potential_energy() 

126 fi2 = self.atoms2.get_forces() 

127 

128 f12[self.mask] += fi1 

129 f12[~self.mask] += fi2 

130 

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

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

133 

134 def get_virtual_charges(self, atoms): 

135 if self.atoms1 is None: 

136 self.initialize(atoms) 

137 

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

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

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

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

142 # TIP4P: OHHX, OHHX, ... 

143 

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

145 vc[self.virtual_mask] = vc1 

146 vc[~self.virtual_mask] = vc2 

147 

148 return vc 

149 

150 def add_virtual_sites(self, positions): 

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

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

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

154 

155 vs[self.virtual_mask] = vs1 

156 vs[~self.virtual_mask] = vs2 

157 

158 return vs 

159 

160 def make_virtual_mask(self): 

161 virtual_mask = [] 

162 ct1 = 0 

163 ct2 = 0 

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

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

166 if self.mask[i]: 

167 ct1 += 1 

168 if not self.mask[i]: 

169 ct2 += 1 

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

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

172 virtual_mask.append(False) 

173 ct2 = 0 

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

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

176 virtual_mask.append(True) 

177 ct1 = 0 

178 

179 self.virtual_mask = np.array(virtual_mask) 

180 

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

182 energy = 0.0 

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

184 

185 self.xpos1 = xpos1 

186 self.xpos2 = xpos2 

187 

188 R1 = xpos1 

189 R2 = xpos2 

190 F1 = np.zeros_like(R1) 

191 F2 = np.zeros_like(R2) 

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

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

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

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

196 # it should not matter much .. 

197 # There is definitely room for improvements here. 

198 cell = self.cell.diagonal() 

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

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

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

202 shift = np.zeros(3) 

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

204 if periodic: 

205 L = cell[i] 

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

207 r00 += shift 

208 

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

210 t = 1 

211 dtdd = 0 

212 if d00 > self.rc: 

213 continue 

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

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

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

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

218 

219 for a1 in range(spm1): 

220 for a2 in range(spm2): 

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

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

223 d = d2**0.5 

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

225 energy += t * e 

226 

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

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

229 

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

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

232 

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

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

235 

236 # Redist forces but dont save forces in org calculators 

237 atoms1 = self.atoms1.copy() 

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

239 atoms1.calc.atoms = atoms1 

240 F1 = atoms1.calc.redistribute_forces(F1) 

241 atoms2 = self.atoms2.copy() 

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

243 atoms2.calc.atoms = atoms2 

244 F2 = atoms2.calc.redistribute_forces(F2) 

245 

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

247 forces[self.mask] = F1 

248 forces[~self.mask] = F2 

249 return energy, forces 

250 

251 def lennard_jones(self, atoms1, atoms2): 

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

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

254 

255 f1 = np.zeros_like(atoms1.positions) 

256 f2 = np.zeros_like(atoms2.positions) 

257 energy = 0.0 

258 

259 cell = self.cell.diagonal() 

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

261 eps = self.epsilon 

262 sig = self.sigma 

263 

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

265 

266 # cutoff from first atom of each mol 

267 shift = np.zeros_like(R00) 

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

269 if periodic: 

270 L = cell[i] 

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

272 R00 += shift 

273 

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

275 d00 = d002**0.5 

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

277 x2 = d00 < self.rc 

278 x12 = np.logical_and(x1, x2) 

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

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

281 t[x2] = 1.0 

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

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

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

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

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

287 continue 

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

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

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

291 c12 = c6**2 

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

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

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

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

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

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

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

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

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

301 

302 return energy, f1, f2 

303 

304 def redistribute_forces(self, forces): 

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

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

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

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

309 f[self.mask] = f1 

310 f[~self.mask] = f2 

311 return f