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