Coverage for ase / calculators / combine_mm.py: 87.92%
207 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-30 08:22 +0000
1# fmt: off
3import copy
5import numpy as np
7from ase import units
8from ase.calculators.calculator import Calculator
9from ase.calculators.qmmm import combine_lj_lorenz_berthelot
11k_c = units.Hartree * units.Bohr
14class CombineMM(Calculator):
15 implemented_properties = ['energy', 'forces']
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, ...)
22 parameters
23 ----------
25 idx: List of indices of atoms belonging to calculator 1
26 apm1,2: atoms pr molecule of each subsystem (NB: apm for TIP4P is 3!)
27 calc1,2: calculator objects for each subsystem
28 sig1,2, eps1,2: LJ parameters for each subsystem. Should be a numpy
29 array of length = apm
30 rc = long range cutoff
31 width = width of cutoff region.
33 Currently the interactions are limited to being:
34 - Nonbonded
35 - Hardcoded to two terms:
36 - Coulomb electrostatics
37 - Lennard-Jones
39 It could of course benefit from being more like the EIQMMM class
40 where the interactions are switchable. But this is in princple
41 just meant for adding counter ions to a qmmm simulation to neutralize
42 the charge of the total systemn
44 Maybe it can combine n MM calculators in the future?
45 """
47 self.idx = idx
48 self.apm1 = apm1 # atoms per mol for LJ calculator
49 self.apm2 = apm2
51 self.rc = rc
52 self.width = width
54 self.atoms1 = None
55 self.atoms2 = None
56 self.mask = None
58 self.calc1 = calc1
59 self.calc2 = calc2
61 self.sig1 = sig1
62 self.eps1 = eps1
63 self.sig2 = sig2
64 self.eps2 = eps2
66 Calculator.__init__(self)
68 def initialize(self, atoms):
69 self.mask = np.zeros(len(atoms), bool)
70 self.mask[self.idx] = True
72 constraints = atoms.constraints
73 atoms.constraints = []
74 self.atoms1 = atoms[self.mask]
75 self.atoms2 = atoms[~self.mask]
77 atoms.constraints = constraints
79 self.atoms1.calc = self.calc1
80 self.atoms2.calc = self.calc2
82 self.cell = atoms.cell
83 self.pbc = atoms.pbc
85 self.sigma, self.epsilon =\
86 combine_lj_lorenz_berthelot(self.sig1, self.sig2,
87 self.eps1, self.eps2)
89 self.make_virtual_mask()
91 def calculate(self, atoms, properties, system_changes):
92 Calculator.calculate(self, atoms, properties, system_changes)
94 if self.atoms1 is None:
95 self.initialize(atoms)
97 pos1 = atoms.positions[self.mask]
98 pos2 = atoms.positions[~self.mask]
99 self.atoms1.set_positions(pos1)
100 self.atoms2.set_positions(pos2)
102 # positions and charges for the coupling term, which should
103 # include virtual charges and sites:
104 spm1 = self.atoms1.calc.sites_per_mol
105 spm2 = self.atoms2.calc.sites_per_mol
106 xpos1 = self.atoms1.calc.add_virtual_sites(pos1)
107 xpos2 = self.atoms2.calc.add_virtual_sites(pos2)
109 xc1 = self.atoms1.calc.get_virtual_charges(self.atoms1)
110 xc2 = self.atoms2.calc.get_virtual_charges(self.atoms2)
112 xpos1 = xpos1.reshape((-1, spm1, 3))
113 xpos2 = xpos2.reshape((-1, spm2, 3))
115 e_c, f_c = self.coulomb(xpos1, xpos2, xc1, xc2, spm1, spm2)
117 e_lj, f1, f2 = self.lennard_jones(self.atoms1, self.atoms2)
119 f_lj = np.zeros((len(atoms), 3))
120 f_lj[self.mask] += f1
121 f_lj[~self.mask] += f2
123 # internal energy, forces of each subsystem:
124 f12 = np.zeros((len(atoms), 3))
125 e1 = self.atoms1.get_potential_energy()
126 fi1 = self.atoms1.get_forces()
128 e2 = self.atoms2.get_potential_energy()
129 fi2 = self.atoms2.get_forces()
131 f12[self.mask] += fi1
132 f12[~self.mask] += fi2
134 self.results['energy'] = e_c + e_lj + e1 + e2
135 self.results['forces'] = f_c + f_lj + f12
137 def get_virtual_charges(self, atoms):
138 if self.atoms1 is None:
139 self.initialize(atoms)
141 vc1 = self.atoms1.calc.get_virtual_charges(atoms[self.mask])
142 vc2 = self.atoms2.calc.get_virtual_charges(atoms[~self.mask])
143 # Need to expand mask with possible new virtual sites.
144 # Virtual sites should ALWAYS be put AFTER actual atoms, like in
145 # TIP4P: OHHX, OHHX, ...
147 vc = np.zeros(len(vc1) + len(vc2))
148 vc[self.virtual_mask] = vc1
149 vc[~self.virtual_mask] = vc2
151 return vc
153 def add_virtual_sites(self, positions):
154 vs1 = self.atoms1.calc.add_virtual_sites(positions[self.mask])
155 vs2 = self.atoms2.calc.add_virtual_sites(positions[~self.mask])
156 vs = np.zeros((len(vs1) + len(vs2), 3))
158 vs[self.virtual_mask] = vs1
159 vs[~self.virtual_mask] = vs2
161 return vs
163 def make_virtual_mask(self):
164 virtual_mask = []
165 ct1 = 0
166 ct2 = 0
167 for i in range(len(self.mask)):
168 virtual_mask.append(self.mask[i])
169 if self.mask[i]:
170 ct1 += 1
171 if not self.mask[i]:
172 ct2 += 1
173 if ((ct2 == self.apm2) &
174 (self.apm2 != self.atoms2.calc.sites_per_mol)):
175 virtual_mask.append(False)
176 ct2 = 0
177 if ((ct1 == self.apm1) &
178 (self.apm1 != self.atoms1.calc.sites_per_mol)):
179 virtual_mask.append(True)
180 ct1 = 0
182 self.virtual_mask = np.array(virtual_mask)
184 def coulomb(self, xpos1, xpos2, xc1, xc2, spm1, spm2):
185 energy = 0.0
186 forces = np.zeros((len(xc1) + len(xc2), 3))
188 self.xpos1 = xpos1
189 self.xpos2 = xpos2
191 R1 = xpos1
192 R2 = xpos2
193 F1 = np.zeros_like(R1)
194 F2 = np.zeros_like(R2)
195 C1 = xc1.reshape((-1, np.shape(xpos1)[1]))
196 C2 = xc2.reshape((-1, np.shape(xpos2)[1]))
197 # Vectorized evaluation is not as trivial when spm1 != spm2.
198 # This is pretty inefficient, but for ~1-5 counter ions as region 1
199 # it should not matter much ..
200 # There is definitely room for improvements here.
201 cell = self.cell.diagonal()
202 for m1, (r1, c1) in enumerate(zip(R1, C1)):
203 for m2, (r2, c2) in enumerate(zip(R2, C2)):
204 r00 = r2[0] - r1[0]
205 shift = np.zeros(3)
206 for i, periodic in enumerate(self.pbc):
207 if periodic:
208 L = cell[i]
209 shift[i] = (r00[i] + L / 2.) % L - L / 2. - r00[i]
210 r00 += shift
212 d00 = (r00**2).sum()**0.5
213 t = 1
214 dtdd = 0
215 if d00 > self.rc:
216 continue
217 elif d00 > self.rc - self.width:
218 y = (d00 - self.rc + self.width) / self.width
219 t -= y**2 * (3.0 - 2.0 * y)
220 dtdd = r00 * 6 * y * (1.0 - y) / (self.width * d00)
222 for a1 in range(spm1):
223 for a2 in range(spm2):
224 r = r2[a2] - r1[a1] + shift
225 d2 = (r**2).sum()
226 d = d2**0.5
227 e = k_c * c1[a1] * c2[a2] / d
228 energy += t * e
230 F1[m1, a1] -= t * (e / d2) * r
231 F2[m2, a2] += t * (e / d2) * r
233 F1[m1, 0] -= dtdd * e
234 F2[m2, 0] += dtdd * e
236 F1 = F1.reshape((-1, 3))
237 F2 = F2.reshape((-1, 3))
239 # Redist forces but dont save forces in org calculators
240 atoms1 = self.atoms1.copy()
241 atoms1.calc = copy.copy(self.calc1)
242 atoms1.calc.atoms = atoms1
243 F1 = atoms1.calc.redistribute_forces(F1)
244 atoms2 = self.atoms2.copy()
245 atoms2.calc = copy.copy(self.calc2)
246 atoms2.calc.atoms = atoms2
247 F2 = atoms2.calc.redistribute_forces(F2)
249 forces = np.zeros((len(self.atoms), 3))
250 forces[self.mask] = F1
251 forces[~self.mask] = F2
252 return energy, forces
254 def lennard_jones(self, atoms1, atoms2):
255 pos1 = atoms1.get_positions().reshape((-1, self.apm1, 3))
256 pos2 = atoms2.get_positions().reshape((-1, self.apm2, 3))
258 f1 = np.zeros_like(atoms1.positions)
259 f2 = np.zeros_like(atoms2.positions)
260 energy = 0.0
262 cell = self.cell.diagonal()
263 for q, p1 in enumerate(pos1): # molwise loop
264 eps = self.epsilon
265 sig = self.sigma
267 R00 = pos2[:, 0] - p1[0, :]
269 # cutoff from first atom of each mol
270 shift = np.zeros_like(R00)
271 for i, periodic in enumerate(self.pbc):
272 if periodic:
273 L = cell[i]
274 shift[:, i] = (R00[:, i] + L / 2) % L - L / 2 - R00[:, i]
275 R00 += shift
277 d002 = (R00**2).sum(1)
278 d00 = d002**0.5
279 x1 = d00 > self.rc - self.width
280 x2 = d00 < self.rc
281 x12 = np.logical_and(x1, x2)
282 y = (d00[x12] - self.rc + self.width) / self.width
283 t = np.zeros(len(d00))
284 t[x2] = 1.0
285 t[x12] -= y**2 * (3.0 - 2.0 * y)
286 dt = np.zeros(len(d00))
287 dt[x12] -= 6.0 / self.width * y * (1.0 - y)
288 for qa in range(len(p1)):
289 if ~np.any(eps[qa, :]):
290 continue
291 R = pos2 - p1[qa, :] + shift[:, None]
292 d2 = (R**2).sum(2)
293 c6 = (sig[qa, :]**2 / d2)**3
294 c12 = c6**2
295 e = 4 * eps[qa, :] * (c12 - c6)
296 energy += np.dot(e.sum(1), t)
297 f = t[:, None, None] * (24 * eps[qa, :] *
298 (2 * c12 - c6) / d2)[:, :, None] * R
299 f00 = - (e.sum(1) * dt / d00)[:, None] * R00
300 f2 += f.reshape((-1, 3))
301 f1[q * self.apm1 + qa, :] -= f.sum(0).sum(0)
302 f1[q * self.apm1, :] -= f00.sum(0)
303 f2[::self.apm2, :] += f00
305 return energy, f1, f2
307 def redistribute_forces(self, forces):
308 f1 = self.calc1.redistribute_forces(forces[self.virtual_mask])
309 f2 = self.calc2.redistribute_forces(forces[~self.virtual_mask])
310 # and then they are back on the real atom centers so
311 f = np.zeros((len(self.atoms), 3))
312 f[self.mask] = f1
313 f[~self.mask] = f2
314 return f