1import numpy as np 

2from ase.calculators.calculator import Calculator 

3from ase.calculators.qmmm import combine_lj_lorenz_berthelot 

4from ase import units 

5import copy 

6 

7k_c = units.Hartree * units.Bohr 

8 

9 

10class CombineMM(Calculator): 

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

12 

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

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

15 """A calculator that combines two MM calculators 

16 (TIPnP, Counterions, ...) 

17 

18 parameters: 

19 

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

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

22 calc1,2: calculator objects for each subsystem 

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

24 array of length = apm 

25 rc = long range cutoff 

26 width = width of cutoff region. 

27 

28 Currently the interactions are limited to being: 

29 - Nonbonded 

30 - Hardcoded to two terms: 

31 - Coulomb electrostatics 

32 - Lennard-Jones 

33 

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

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

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

37 the charge of the total systemn 

38 

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

40 """ 

41 

42 self.idx = idx 

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

44 self.apm2 = apm2 

45 

46 self.rc = rc 

47 self.width = width 

48 

49 self.atoms1 = None 

50 self.atoms2 = None 

51 self.mask = None 

52 

53 self.calc1 = calc1 

54 self.calc2 = calc2 

55 

56 self.sig1 = sig1 

57 self.eps1 = eps1 

58 self.sig2 = sig2 

59 self.eps2 = eps2 

60 

61 Calculator.__init__(self) 

62 

63 def initialize(self, atoms): 

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

65 self.mask[self.idx] = True 

66 

67 constraints = atoms.constraints 

68 atoms.constraints = [] 

69 self.atoms1 = atoms[self.mask] 

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

71 

72 atoms.constraints = constraints 

73 

74 self.atoms1.calc = self.calc1 

75 self.atoms2.calc = self.calc2 

76 

77 self.cell = atoms.cell 

78 self.pbc = atoms.pbc 

79 

80 self.sigma, self.epsilon =\ 

81 combine_lj_lorenz_berthelot(self.sig1, self.sig2, 

82 self.eps1, self.eps2) 

83 

84 self.make_virtual_mask() 

85 

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

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

88 

89 if self.atoms1 is None: 

90 self.initialize(atoms) 

91 

92 pos1 = atoms.positions[self.mask] 

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

94 self.atoms1.set_positions(pos1) 

95 self.atoms2.set_positions(pos2) 

96 

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

98 # include virtual charges and sites: 

99 spm1 = self.atoms1.calc.sites_per_mol 

100 spm2 = self.atoms2.calc.sites_per_mol 

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

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

103 

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

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

106 

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

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

109 

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

111 

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

113 

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

115 f_lj[self.mask] += f1 

116 f_lj[~self.mask] += f2 

117 

118 # internal energy, forces of each subsystem: 

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

120 e1 = self.atoms1.get_potential_energy() 

121 fi1 = self.atoms1.get_forces() 

122 

123 e2 = self.atoms2.get_potential_energy() 

124 fi2 = self.atoms2.get_forces() 

125 

126 f12[self.mask] += fi1 

127 f12[~self.mask] += fi2 

128 

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

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

131 

132 def get_virtual_charges(self, atoms): 

133 if self.atoms1 is None: 

134 self.initialize(atoms) 

135 

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

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

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

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

140 # TIP4P: OHHX, OHHX, ... 

141 

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

143 vc[self.virtual_mask] = vc1 

144 vc[~self.virtual_mask] = vc2 

145 

146 return vc 

147 

148 def add_virtual_sites(self, positions): 

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

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

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

152 

153 vs[self.virtual_mask] = vs1 

154 vs[~self.virtual_mask] = vs2 

155 

156 return vs 

157 

158 def make_virtual_mask(self): 

159 virtual_mask = [] 

160 ct1 = 0 

161 ct2 = 0 

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

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

164 if self.mask[i]: 

165 ct1 += 1 

166 if not self.mask[i]: 

167 ct2 += 1 

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

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

170 virtual_mask.append(False) 

171 ct2 = 0 

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

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

174 virtual_mask.append(True) 

175 ct1 = 0 

176 

177 self.virtual_mask = np.array(virtual_mask) 

178 

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

180 energy = 0.0 

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

182 

183 self.xpos1 = xpos1 

184 self.xpos2 = xpos2 

185 

186 R1 = xpos1 

187 R2 = xpos2 

188 F1 = np.zeros_like(R1) 

189 F2 = np.zeros_like(R2) 

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

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

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

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

194 # it should not matter much .. 

195 # There is definitely room for improvements here. 

196 cell = self.cell.diagonal() 

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

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

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

200 shift = np.zeros(3) 

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

202 if periodic: 

203 L = cell[i] 

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

205 r00 += shift 

206 

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

208 t = 1 

209 dtdd = 0 

210 if d00 > self.rc: 

211 continue 

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

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

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

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

216 

217 for a1 in range(spm1): 

218 for a2 in range(spm2): 

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

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

221 d = d2**0.5 

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

223 energy += t * e 

224 

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

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

227 

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

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

230 

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

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

233 

234 # Redist forces but dont save forces in org calculators 

235 atoms1 = self.atoms1.copy() 

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

237 atoms1.calc.atoms = atoms1 

238 F1 = atoms1.calc.redistribute_forces(F1) 

239 atoms2 = self.atoms2.copy() 

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

241 atoms2.calc.atoms = atoms2 

242 F2 = atoms2.calc.redistribute_forces(F2) 

243 

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

245 forces[self.mask] = F1 

246 forces[~self.mask] = F2 

247 return energy, forces 

248 

249 def lennard_jones(self, atoms1, atoms2): 

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

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

252 

253 f1 = np.zeros_like(atoms1.positions) 

254 f2 = np.zeros_like(atoms2.positions) 

255 energy = 0.0 

256 

257 cell = self.cell.diagonal() 

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

259 eps = self.epsilon 

260 sig = self.sigma 

261 

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

263 

264 # cutoff from first atom of each mol 

265 shift = np.zeros_like(R00) 

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

267 if periodic: 

268 L = cell[i] 

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

270 R00 += shift 

271 

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

273 d00 = d002**0.5 

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

275 x2 = d00 < self.rc 

276 x12 = np.logical_and(x1, x2) 

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

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

279 t[x2] = 1.0 

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

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

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

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

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

285 continue 

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

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

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

289 c12 = c6**2 

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

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

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

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

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

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

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

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

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

299 

300 return energy, f1, f2 

301 

302 def redistribute_forces(self, forces): 

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

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

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

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

307 f[self.mask] = f1 

308 f[~self.mask] = f2 

309 return f