Coverage for ase / calculators / qmmm.py: 90.99%

455 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-03-30 08:22 +0000

1# fmt: off 

2 

3from collections.abc import Sequence 

4 

5import numpy as np 

6 

7from ase.calculators.calculator import Calculator 

8from ase.cell import Cell 

9from ase.data import atomic_numbers 

10from ase.geometry import get_distances 

11from ase.parallel import world 

12from ase.utils import IOContext 

13 

14 

15class SimpleQMMM(Calculator): 

16 """Simple QMMM calculator.""" 

17 

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

19 

20 def __init__(self, selection, qmcalc, mmcalc1, mmcalc2, vacuum=None): 

21 """SimpleQMMM object. 

22 

23 The energy is calculated as:: 

24 

25 _ _ _ 

26 E = E (R ) - E (R ) + E (R ) 

27 QM QM MM QM MM all 

28 

29 parameters 

30 ---------- 

31 

32 selection: list of int, slice object or list of bool 

33 Selection out of all the atoms that belong to the QM part. 

34 qmcalc: Calculator object 

35 QM-calculator. 

36 mmcalc1: Calculator object 

37 MM-calculator used for QM region. 

38 mmcalc2: Calculator object 

39 MM-calculator used for everything. 

40 vacuum: float or None 

41 Amount of vacuum to add around QM atoms. Use None if QM 

42 calculator doesn't need a box. 

43 

44 """ 

45 self.selection = selection 

46 self.qmcalc = qmcalc 

47 self.mmcalc1 = mmcalc1 

48 self.mmcalc2 = mmcalc2 

49 self.vacuum = vacuum 

50 

51 self.qmatoms = None 

52 self.center = None 

53 

54 Calculator.__init__(self) 

55 

56 def _get_name(self): 

57 return f'{self.qmcalc.name}-{self.mmcalc1.name}+{self.mmcalc1.name}' 

58 

59 def initialize_qm(self, atoms): 

60 constraints = atoms.constraints 

61 atoms.constraints = [] 

62 self.qmatoms = atoms[self.selection] 

63 atoms.constraints = constraints 

64 self.qmatoms.pbc = False 

65 if self.vacuum: 

66 self.qmatoms.center(vacuum=self.vacuum) 

67 self.center = self.qmatoms.positions.mean(axis=0) 

68 

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

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

71 

72 if self.qmatoms is None: 

73 self.initialize_qm(atoms) 

74 

75 self.qmatoms.positions = atoms.positions[self.selection] 

76 if self.vacuum: 

77 self.qmatoms.positions += (self.center - 

78 self.qmatoms.positions.mean(axis=0)) 

79 

80 energy = self.qmcalc.get_potential_energy(self.qmatoms) 

81 qmforces = self.qmcalc.get_forces(self.qmatoms) 

82 energy += self.mmcalc2.get_potential_energy(atoms) 

83 forces = self.mmcalc2.get_forces(atoms) 

84 

85 if self.vacuum: 

86 qmforces -= qmforces.mean(axis=0) 

87 forces[self.selection] += qmforces 

88 

89 energy -= self.mmcalc1.get_potential_energy(self.qmatoms) 

90 forces[self.selection] -= self.mmcalc1.get_forces(self.qmatoms) 

91 

92 self.results['energy'] = energy 

93 self.results['forces'] = forces 

94 

95 

96class EIQMMM(Calculator, IOContext): 

97 """Explicit interaction QMMM calculator.""" 

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

99 

100 def __init__(self, selection, qmcalc, mmcalc, interaction, 

101 vacuum=None, embedding=None, output=None, comm=world): 

102 """EIQMMM object. 

103 

104 The energy is calculated as:: 

105 

106 _ _ _ _ 

107 E = E (R ) + E (R ) + E (R , R ) 

108 QM QM MM MM I QM MM 

109 

110 parameters 

111 ---------- 

112 

113 selection: list of int, slice object or list of bool 

114 Selection out of all the atoms that belong to the QM part. 

115 qmcalc: Calculator object 

116 QM-calculator. 

117 mmcalc: Calculator object 

118 MM-calculator. 

119 interaction: Interaction object 

120 Interaction between QM and MM regions. 

121 vacuum: float or None 

122 Amount of vacuum to add around QM atoms. Use None if QM 

123 calculator doesn't need a box. 

124 embedding: Embedding object or None 

125 Specialized embedding object. Use None in order to use the 

126 default one. 

127 output: None, '-', str or file-descriptor. 

128 File for logging information - default is no logging (None). 

129 

130 """ 

131 

132 self.selection = selection 

133 

134 self.qmcalc = qmcalc 

135 self.mmcalc = mmcalc 

136 self.interaction = interaction 

137 self.vacuum = vacuum 

138 self.embedding = embedding 

139 

140 self.qmatoms = None 

141 self.mmatoms = None 

142 self.mask = None 

143 self.center = None # center of QM atoms in QM-box 

144 

145 self.output = self.openfile(file=output, comm=comm) 

146 

147 Calculator.__init__(self) 

148 

149 def _get_name(self): 

150 return f'{self.qmcalc.name}+{self.interaction.name}+{self.mmcalc.name}' 

151 

152 def initialize(self, atoms): 

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

154 self.mask[self.selection] = True 

155 

156 constraints = atoms.constraints 

157 atoms.constraints = [] # avoid slicing of constraints 

158 self.qmatoms = atoms[self.mask] 

159 self.mmatoms = atoms[~self.mask] 

160 atoms.constraints = constraints 

161 

162 self.qmatoms.pbc = False 

163 

164 if self.vacuum: 

165 self.qmatoms.center(vacuum=self.vacuum) 

166 self.center = self.qmatoms.positions.mean(axis=0) 

167 print('Size of QM-cell after centering:', 

168 self.qmatoms.cell.diagonal(), file=self.output) 

169 

170 self.qmatoms.calc = self.qmcalc 

171 self.mmatoms.calc = self.mmcalc 

172 

173 if self.embedding is None: 

174 self.embedding = Embedding() 

175 

176 self.embedding.initialize(self.qmatoms, self.mmatoms) 

177 print('Embedding:', self.embedding, file=self.output) 

178 

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

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

181 

182 if self.qmatoms is None: 

183 self.initialize(atoms) 

184 

185 self.mmatoms.set_positions(atoms.positions[~self.mask]) 

186 self.qmatoms.set_positions(atoms.positions[self.mask]) 

187 

188 if self.vacuum: 

189 shift = self.center - self.qmatoms.positions.mean(axis=0) 

190 self.qmatoms.positions += shift 

191 else: 

192 shift = (0, 0, 0) 

193 

194 self.embedding.update(shift) 

195 

196 ienergy, iqmforces, immforces = self.interaction.calculate( 

197 self.qmatoms, self.mmatoms, shift) 

198 

199 qmenergy = self.qmatoms.get_potential_energy() 

200 mmenergy = self.mmatoms.get_potential_energy() 

201 energy = ienergy + qmenergy + mmenergy 

202 

203 print('Energies: {:12.3f} {:+12.3f} {:+12.3f} = {:12.3f}' 

204 .format(ienergy, qmenergy, mmenergy, energy), file=self.output) 

205 

206 qmforces = self.qmatoms.get_forces() 

207 mmforces = self.mmatoms.get_forces() 

208 

209 mmforces += self.embedding.get_mm_forces() 

210 

211 forces = np.empty((len(atoms), 3)) 

212 forces[self.mask] = qmforces + iqmforces 

213 forces[~self.mask] = mmforces + immforces 

214 

215 self.results['energy'] = energy 

216 self.results['forces'] = forces 

217 

218 

219def wrap(D, cell, pbc): 

220 """Wrap distances to nearest neighbor (minimum image convention).""" 

221 for i, periodic in enumerate(pbc): 

222 if periodic: 

223 d = D[:, i] 

224 L = cell[i] 

225 d[:] = (d + L / 2) % L - L / 2 # modify D inplace 

226 

227 

228class Embedding: 

229 def __init__(self, molecule_size=3, **parameters): 

230 """Point-charge embedding.""" 

231 self.qmatoms = None 

232 self.mmatoms = None 

233 self.molecule_size = molecule_size 

234 self.virtual_molecule_size = None 

235 self.parameters = parameters 

236 

237 def __repr__(self): 

238 return f'Embedding(molecule_size={self.molecule_size})' 

239 

240 def initialize(self, qmatoms, mmatoms): 

241 """Hook up embedding object to QM and MM atoms objects.""" 

242 self.qmatoms = qmatoms 

243 self.mmatoms = mmatoms 

244 charges = mmatoms.calc.get_virtual_charges(mmatoms) 

245 self.pcpot = qmatoms.calc.embed(charges, **self.parameters) 

246 self.virtual_molecule_size = (self.molecule_size * 

247 len(charges) // len(mmatoms)) 

248 

249 def update(self, shift): 

250 """Update point-charge positions.""" 

251 # Wrap point-charge positions to the MM-cell closest to the 

252 # center of the the QM box, but avoid ripping molecules apart: 

253 qmcenter = self.qmatoms.positions.mean(axis=0) 

254 # if counter ions are used, then molecule_size has more than 1 value 

255 if self.mmatoms.calc.name == 'combinemm': 

256 mask1 = self.mmatoms.calc.mask 

257 mask2 = ~mask1 

258 vmask1 = self.mmatoms.calc.virtual_mask 

259 vmask2 = ~vmask1 

260 apm1 = self.mmatoms.calc.apm1 

261 apm2 = self.mmatoms.calc.apm2 

262 spm1 = self.mmatoms.calc.atoms1.calc.sites_per_mol 

263 spm2 = self.mmatoms.calc.atoms2.calc.sites_per_mol 

264 pos = self.mmatoms.positions 

265 pos1 = pos[mask1].reshape((-1, apm1, 3)) 

266 pos2 = pos[mask2].reshape((-1, apm2, 3)) 

267 pos = (pos1, pos2) 

268 else: 

269 pos = (self.mmatoms.positions, ) 

270 apm1 = self.molecule_size 

271 apm2 = self.molecule_size 

272 # This is only specific to calculators where apm != spm, 

273 # i.e. TIP4P. Non-native MM calcs do not have this attr. 

274 if hasattr(self.mmatoms.calc, 'sites_per_mol'): 

275 spm1 = self.mmatoms.calc.sites_per_mol 

276 spm2 = self.mmatoms.calc.sites_per_mol 

277 else: 

278 spm1 = self.molecule_size 

279 spm2 = spm1 

280 mask1 = np.ones(len(self.mmatoms), dtype=bool) 

281 mask2 = mask1 

282 

283 wrap_pos = np.zeros_like(self.mmatoms.positions) 

284 com_all = [] 

285 apm = (apm1, apm2) 

286 mask = (mask1, mask2) 

287 spm = (spm1, spm2) 

288 for p, n, m, vn in zip(pos, apm, mask, spm): 

289 positions = p.reshape((-1, n, 3)) + shift 

290 

291 # Distances from the center of the QM box to the first atom of 

292 # each molecule: 

293 distances = positions[:, 0] - qmcenter 

294 

295 wrap(distances, self.mmatoms.cell.diagonal(), self.mmatoms.pbc) 

296 offsets = distances - positions[:, 0] 

297 positions += offsets[:, np.newaxis] + qmcenter 

298 

299 # Geometric center positions for each mm mol for LR cut 

300 com = np.array([p.mean(axis=0) for p in positions]) 

301 # Need per atom for C-code: 

302 com_pv = np.repeat(com, vn, axis=0) 

303 com_all.append(com_pv) 

304 

305 wrap_pos[m] = positions.reshape((-1, 3)) 

306 

307 positions = wrap_pos.copy() 

308 positions = self.mmatoms.calc.add_virtual_sites(positions) 

309 

310 if self.mmatoms.calc.name == 'combinemm': 

311 com_pv = np.zeros_like(positions) 

312 for ii, m in enumerate((vmask1, vmask2)): 

313 com_pv[m] = com_all[ii] 

314 

315 # compatibility with gpaw versions w/o LR cut in PointChargePotential 

316 if 'rc2' in self.parameters: 

317 self.pcpot.set_positions(positions, com_pv=com_pv) 

318 else: 

319 self.pcpot.set_positions(positions) 

320 

321 def get_mm_forces(self): 

322 """Calculate the forces on the MM-atoms from the QM-part.""" 

323 f = self.pcpot.get_forces(self.qmatoms.calc) 

324 return self.mmatoms.calc.redistribute_forces(f) 

325 

326 

327def combine_lj_lorenz_berthelot(sigmaqm, sigmamm, 

328 epsilonqm, epsilonmm): 

329 """Combine LJ parameters according to the Lorenz-Berthelot rule""" 

330 sigma = [] 

331 epsilon = [] 

332 # check if input is tuple of vals for more than 1 mm calc, or only for 1. 

333 if isinstance(sigmamm, Sequence): 

334 numcalcs = len(sigmamm) 

335 else: 

336 numcalcs = 1 # if only 1 mm calc, eps and sig are simply np arrays 

337 sigmamm = (sigmamm, ) 

338 epsilonmm = (epsilonmm, ) 

339 for cc in range(numcalcs): 

340 sigma_c = np.zeros((len(sigmaqm), len(sigmamm[cc]))) 

341 epsilon_c = np.zeros_like(sigma_c) 

342 

343 for ii in range(len(sigmaqm)): 

344 sigma_c[ii, :] = (sigmaqm[ii] + sigmamm[cc]) / 2 

345 epsilon_c[ii, :] = (epsilonqm[ii] * epsilonmm[cc])**0.5 

346 sigma.append(sigma_c) 

347 epsilon.append(epsilon_c) 

348 

349 if numcalcs == 1: # retain original, 1 calc function 

350 sigma = np.array(sigma[0]) 

351 epsilon = np.array(epsilon[0]) 

352 

353 return sigma, epsilon 

354 

355 

356class LJInteractionsGeneral: 

357 name = 'LJ-general' 

358 

359 def __init__(self, sigmaqm, epsilonqm, sigmamm, epsilonmm, 

360 qm_molecule_size, mm_molecule_size=3, 

361 rc=np.inf, width=1.0): 

362 """General Lennard-Jones type explicit interaction. 

363 

364 sigmaqm: array 

365 Array of sigma-parameters which should have the length of the QM 

366 subsystem 

367 epsilonqm: array 

368 As sigmaqm, but for epsilon-paramaters 

369 sigmamm: Either array (A) or tuple (B) 

370 A (no counterions): 

371 Array of sigma-parameters with the length of the smallests 

372 repeating atoms-group (i.e. molecule) of the MM subsystem 

373 B (counterions): 

374 Tuple: (arr1, arr2), where arr1 is an array of sigmas with 

375 the length of counterions in the MM subsystem, and 

376 arr2 is the array from A. 

377 epsilonmm: array or tuple 

378 As sigmamm but for epsilon-parameters. 

379 qm_molecule_size: int 

380 number of atoms of the smallest repeating atoms-group (i.e. 

381 molecule) in the QM subsystem (often just the number of atoms 

382 in the QM subsystem) 

383 mm_molecule_size: int 

384 as qm_molecule_size but for the MM subsystem. Will be overwritten 

385 if counterions are present in the MM subsystem (via the CombineMM 

386 calculator) 

387 

388 """ 

389 self.sigmaqm = sigmaqm 

390 self.epsilonqm = epsilonqm 

391 self.sigmamm = sigmamm 

392 self.epsilonmm = epsilonmm 

393 self.qms = qm_molecule_size 

394 self.mms = mm_molecule_size 

395 self.rc = rc 

396 self.width = width 

397 self.combine_lj() 

398 

399 def combine_lj(self): 

400 self.sigma, self.epsilon = combine_lj_lorenz_berthelot( 

401 self.sigmaqm, self.sigmamm, self.epsilonqm, self.epsilonmm) 

402 

403 def calculate(self, qmatoms, mmatoms, shift): 

404 epsilon = self.epsilon 

405 sigma = self.sigma 

406 

407 # loop over possible multiple mm calculators 

408 # currently 1 or 2, but could be generalized in the future... 

409 apm1 = self.mms 

410 mask1 = np.ones(len(mmatoms), dtype=bool) 

411 mask2 = mask1 

412 apm = (apm1, ) 

413 sigma = (sigma, ) 

414 epsilon = (epsilon, ) 

415 if hasattr(mmatoms.calc, 'name'): 

416 if mmatoms.calc.name == 'combinemm': 

417 mask1 = mmatoms.calc.mask 

418 mask2 = ~mask1 

419 apm1 = mmatoms.calc.apm1 

420 apm2 = mmatoms.calc.apm2 

421 apm = (apm1, apm2) 

422 sigma = sigma[0] # Was already loopable 2-tuple 

423 epsilon = epsilon[0] 

424 

425 mask = (mask1, mask2) 

426 e_all = 0 

427 qmforces_all = np.zeros_like(qmatoms.positions) 

428 mmforces_all = np.zeros_like(mmatoms.positions) 

429 

430 # zip stops at shortest tuple so we dont double count 

431 # cases of no counter ions. 

432 for n, m, eps, sig in zip(apm, mask, epsilon, sigma): 

433 mmpositions = self.update(qmatoms, mmatoms[m], n, shift) 

434 qmforces = np.zeros_like(qmatoms.positions) 

435 mmforces = np.zeros_like(mmatoms[m].positions) 

436 energy = 0.0 

437 

438 qmpositions = qmatoms.positions.reshape((-1, self.qms, 3)) 

439 

440 for q, qmpos in enumerate(qmpositions): # molwise loop 

441 # cutoff from first atom of each mol 

442 R00 = mmpositions[:, 0] - qmpos[0, :] 

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

444 d00 = d002**0.5 

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

446 x2 = d00 < self.rc 

447 x12 = np.logical_and(x1, x2) 

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

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

450 t[x2] = 1.0 

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

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

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

454 for qa in range(len(qmpos)): 

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

456 continue 

457 R = mmpositions - qmpos[qa, :] 

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

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

460 c12 = c6**2 

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

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

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

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

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

466 mmforces += f.reshape((-1, 3)) 

467 qmforces[q * self.qms + qa, :] -= f.sum(0).sum(0) 

468 qmforces[q * self.qms, :] -= f00.sum(0) 

469 mmforces[::n, :] += f00 

470 

471 e_all += energy 

472 qmforces_all += qmforces 

473 mmforces_all[m] += mmforces 

474 

475 return e_all, qmforces_all, mmforces_all 

476 

477 def update(self, qmatoms, mmatoms, n, shift): 

478 # Wrap point-charge positions to the MM-cell closest to the 

479 # center of the the QM box, but avoid ripping molecules apart: 

480 qmcenter = qmatoms.cell.diagonal() / 2 

481 positions = mmatoms.positions.reshape((-1, n, 3)) + shift 

482 

483 # Distances from the center of the QM box to the first atom of 

484 # each molecule: 

485 distances = positions[:, 0] - qmcenter 

486 

487 wrap(distances, mmatoms.cell.diagonal(), mmatoms.pbc) 

488 offsets = distances - positions[:, 0] 

489 positions += offsets[:, np.newaxis] + qmcenter 

490 

491 return positions 

492 

493 

494class LJInteractions: 

495 name = 'LJ' 

496 

497 def __init__(self, parameters): 

498 """Lennard-Jones type explicit interaction. 

499 

500 parameters: dict 

501 Mapping from pair of atoms to tuple containing epsilon and sigma 

502 for that pair. 

503 

504 Example: 

505 

506 lj = LJInteractions({('O', 'O'): (eps, sigma)}) 

507 

508 """ 

509 self.parameters = {} 

510 for (symbol1, symbol2), (epsilon, sigma) in parameters.items(): 

511 Z1 = atomic_numbers[symbol1] 

512 Z2 = atomic_numbers[symbol2] 

513 self.parameters[(Z1, Z2)] = epsilon, sigma 

514 self.parameters[(Z2, Z1)] = epsilon, sigma 

515 

516 def calculate(self, qmatoms, mmatoms, shift): 

517 qmforces = np.zeros_like(qmatoms.positions) 

518 mmforces = np.zeros_like(mmatoms.positions) 

519 species = set(mmatoms.numbers) 

520 energy = 0.0 

521 for R1, Z1, F1 in zip(qmatoms.positions, qmatoms.numbers, qmforces): 

522 for Z2 in species: 

523 if (Z1, Z2) not in self.parameters: 

524 continue 

525 epsilon, sigma = self.parameters[(Z1, Z2)] 

526 mask = (mmatoms.numbers == Z2) 

527 D = mmatoms.positions[mask] + shift - R1 

528 wrap(D, mmatoms.cell.diagonal(), mmatoms.pbc) 

529 d2 = (D**2).sum(1) 

530 c6 = (sigma**2 / d2)**3 

531 c12 = c6**2 

532 energy += 4 * epsilon * (c12 - c6).sum() 

533 f = 24 * epsilon * ((2 * c12 - c6) / d2)[:, np.newaxis] * D 

534 F1 -= f.sum(0) 

535 mmforces[mask] += f 

536 return energy, qmforces, mmforces 

537 

538 

539class RescaledCalculator(Calculator): 

540 """Rescales length and energy of a calculators to match given 

541 lattice constant and bulk modulus 

542 

543 Useful for MM calculator used within a :class:`ForceQMMM` model. 

544 See T. D. Swinburne and J. R. Kermode, Phys. Rev. B 96, 144102 (2017) 

545 for a derivation of the scaling constants. 

546 """ 

547 implemented_properties = ['forces', 'energy', 'stress'] 

548 

549 def __init__(self, mm_calc, 

550 qm_lattice_constant, qm_bulk_modulus, 

551 mm_lattice_constant, mm_bulk_modulus): 

552 Calculator.__init__(self) 

553 self.mm_calc = mm_calc 

554 self.alpha = qm_lattice_constant / mm_lattice_constant 

555 self.beta = mm_bulk_modulus / qm_bulk_modulus / (self.alpha**3) 

556 

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

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

559 

560 # mm_pos = atoms.get_positions() 

561 scaled_atoms = atoms.copy() 

562 

563 # scaled_atoms.positions = mm_pos/self.alpha 

564 mm_cell = atoms.get_cell() 

565 scaled_atoms.set_cell(mm_cell / self.alpha, scale_atoms=True) 

566 

567 results = {} 

568 

569 if 'energy' in properties: 

570 energy = self.mm_calc.get_potential_energy(scaled_atoms) 

571 results['energy'] = energy / self.beta 

572 

573 if 'forces' in properties: 

574 forces = self.mm_calc.get_forces(scaled_atoms) 

575 results['forces'] = forces / (self.beta * self.alpha) 

576 

577 if 'stress' in properties: 

578 stress = self.mm_calc.get_stress(scaled_atoms) 

579 results['stress'] = stress / (self.beta * self.alpha**3) 

580 

581 self.results = results 

582 

583 

584class ForceConstantCalculator(Calculator): 

585 """ 

586 Compute forces based on provided force-constant matrix 

587 

588 Useful with `ForceQMMM` to do harmonic QM/MM using force constants 

589 of QM method. 

590 """ 

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

592 

593 def __init__(self, D, ref, f0): 

594 """ 

595 Parameters 

596 ---------- 

597 

598 D: matrix or sparse matrix, shape `(3*len(ref), 3*len(ref))` 

599 Force constant matrix. 

600 Sign convention is `D_ij = d^2E/(dx_i dx_j), so 

601 `force = -D.dot(displacement)` 

602 ref: ase.atoms.Atoms 

603 Atoms object for reference configuration 

604 f0: array, shape `(len(ref), 3)` 

605 Value of forces at reference configuration 

606 """ 

607 assert D.shape[0] == D.shape[1] 

608 assert D.shape[0] // 3 == len(ref) 

609 self.D = D 

610 self.ref = ref 

611 self.f0 = f0 

612 self.size = len(ref) 

613 Calculator.__init__(self) 

614 

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

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

617 u = atoms.positions - self.ref.positions 

618 f = -self.D.dot(u.reshape(3 * self.size)) 

619 forces = np.zeros((len(atoms), 3)) 

620 forces[:, :] = f.reshape(self.size, 3) 

621 self.results['forces'] = forces + self.f0 

622 self.results['energy'] = 0.0 

623 

624 

625class ForceQMMM(Calculator): 

626 """ 

627 Force-based QM/MM calculator 

628 

629 QM forces are computed using a buffer region and then mixed abruptly 

630 with MM forces: 

631 

632 F^i_QMMM = { F^i_QM if i in QM region 

633 { F^i_MM otherwise 

634 

635 cf. N. Bernstein, J. R. Kermode, and G. Csanyi, 

636 Rep. Prog. Phys. 72, 026501 (2009) 

637 and T. D. Swinburne and J. R. Kermode, Phys. Rev. B 96, 144102 (2017). 

638 """ 

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

640 

641 def __init__(self, 

642 atoms, 

643 qm_selection_mask, 

644 qm_calc, 

645 mm_calc, 

646 buffer_width, 

647 vacuum=5., 

648 zero_mean=True, 

649 qm_cell_round_off=3, 

650 qm_radius=None): 

651 """ 

652 ForceQMMM calculator 

653 

654 Parameters 

655 ---------- 

656 

657 qm_selection_mask: list of ints, slice object or bool list/array 

658 Selection out of atoms that belong to the QM region. 

659 qm_calc: Calculator object 

660 QM-calculator. 

661 mm_calc: Calculator object 

662 MM-calculator (should be scaled, see :class:`RescaledCalculator`) 

663 Can use `ForceConstantCalculator` based on QM force constants, if 

664 available. 

665 vacuum: float or None 

666 Amount of vacuum to add around QM atoms. 

667 zero_mean: bool 

668 If True, add a correction to zero the mean force in each direction 

669 qm_cell_round_off: float 

670 Tolerance value in Angstrom to round the qm cluster cell 

671 qm_radius: 3x1 array of floats qm_radius for [x, y, z] 

672 3d qm radius for calculation of qm cluster cell. default is None 

673 and the radius is estimated from maximum distance between the atoms 

674 in qm region. 

675 """ 

676 

677 if len(atoms[qm_selection_mask]) == 0: 

678 raise ValueError("no QM atoms selected!") 

679 

680 self.qm_selection_mask = qm_selection_mask 

681 self.qm_calc = qm_calc 

682 self.mm_calc = mm_calc 

683 self.vacuum = vacuum 

684 self.buffer_width = buffer_width 

685 self.zero_mean = zero_mean 

686 self.qm_cell_round_off = qm_cell_round_off 

687 self.qm_radius = qm_radius 

688 

689 self.qm_buffer_mask = None 

690 

691 Calculator.__init__(self) 

692 

693 def initialize_qm_buffer_mask(self, atoms): 

694 """ 

695 Initialises system to perform qm calculation 

696 """ 

697 # calculate the distances between all atoms and qm atoms 

698 # qm_distance_matrix is a [N_QM_atoms x N_atoms] matrix 

699 _, qm_distance_matrix = get_distances( 

700 atoms.positions[self.qm_selection_mask], 

701 atoms.positions, 

702 atoms.cell, atoms.pbc) 

703 

704 self.qm_buffer_mask = np.zeros(len(atoms), dtype=bool) 

705 

706 # every r_qm is a matrix of distances 

707 # from an atom in qm region and all atoms with size [N_atoms] 

708 for r_qm in qm_distance_matrix: 

709 self.qm_buffer_mask[r_qm < self.buffer_width] = True 

710 

711 def get_qm_cluster(self, atoms): 

712 

713 if self.qm_buffer_mask is None: 

714 self.initialize_qm_buffer_mask(atoms) 

715 

716 qm_cluster = atoms[self.qm_buffer_mask] 

717 del qm_cluster.constraints 

718 

719 round_cell = False 

720 if self.qm_radius is None: 

721 round_cell = True 

722 # get all distances between qm atoms. 

723 # Treat all X, Y and Z directions independently 

724 # only distance between qm atoms is calculated 

725 # in order to estimate qm radius in thee directions 

726 R_qm, _ = get_distances(atoms.positions[self.qm_selection_mask], 

727 cell=atoms.cell, pbc=atoms.pbc) 

728 # estimate qm radius in three directions as 1/2 

729 # of max distance between qm atoms 

730 self.qm_radius = np.amax(np.amax(R_qm, axis=1), axis=0) * 0.5 

731 

732 if atoms.cell.orthorhombic: 

733 cell_size = np.diagonal(atoms.cell) 

734 else: 

735 raise RuntimeError("NON-orthorhombic cell is not supported!") 

736 

737 # check if qm_cluster should be left periodic 

738 # in periodic directions of the cell (cell[i] < qm_radius + buffer 

739 # otherwise change to non pbc 

740 # and make a cluster in a vacuum configuration 

741 qm_cluster_pbc = (atoms.pbc & 

742 (cell_size < 

743 2.0 * (self.qm_radius + self.buffer_width))) 

744 

745 # start with the original orthorhombic cell 

746 qm_cluster_cell = cell_size.copy() 

747 # create a cluster in a vacuum cell in non periodic directions 

748 qm_cluster_cell[~qm_cluster_pbc] = ( 

749 2.0 * (self.qm_radius[~qm_cluster_pbc] + 

750 self.buffer_width + 

751 self.vacuum)) 

752 

753 if round_cell: 

754 # round the qm cell to the required tolerance 

755 qm_cluster_cell[~qm_cluster_pbc] = (np.round( 

756 (qm_cluster_cell[~qm_cluster_pbc]) / 

757 self.qm_cell_round_off) * 

758 self.qm_cell_round_off) 

759 

760 qm_cluster.set_cell(Cell(np.diag(qm_cluster_cell))) 

761 qm_cluster.pbc = qm_cluster_pbc 

762 

763 qm_shift = (0.5 * qm_cluster.cell.diagonal() - 

764 qm_cluster.positions.mean(axis=0)) 

765 

766 if 'cell_origin' in qm_cluster.info: 

767 del qm_cluster.info['cell_origin'] 

768 

769 # center the cluster only in non pbc directions 

770 qm_cluster.positions[:, ~qm_cluster_pbc] += qm_shift[~qm_cluster_pbc] 

771 

772 return qm_cluster 

773 

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

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

776 

777 qm_cluster = self.get_qm_cluster(atoms) 

778 

779 forces = self.mm_calc.get_forces(atoms) 

780 qm_forces = self.qm_calc.get_forces(qm_cluster) 

781 

782 forces[self.qm_selection_mask] = \ 

783 qm_forces[self.qm_selection_mask[self.qm_buffer_mask]] 

784 

785 if self.zero_mean: 

786 # Target is that: forces.sum(axis=1) == [0., 0., 0.] 

787 forces[:] -= forces.mean(axis=0) 

788 

789 self.results['forces'] = forces 

790 self.results['energy'] = 0.0 

791 

792 def get_region_from_masks(self, atoms=None, print_mapping=False): 

793 """ 

794 creates region array from the masks of the calculators. The tags in 

795 the array are: 

796 QM - qm atoms 

797 buffer - buffer atoms 

798 MM - atoms treated with mm calculator 

799 """ 

800 if atoms is None: 

801 if self.atoms is None: 

802 raise ValueError('Calculator has no atoms') 

803 else: 

804 atoms = self.atoms 

805 

806 region = np.full_like(atoms, "MM") 

807 

808 region[self.qm_selection_mask] = ( 

809 np.full_like(region[self.qm_selection_mask], "QM")) 

810 

811 buffer_only_mask = self.qm_buffer_mask & ~self.qm_selection_mask 

812 

813 region[buffer_only_mask] = np.full_like(region[buffer_only_mask], 

814 "buffer") 

815 

816 if print_mapping: 

817 

818 print(f"Mapping of {len(region):5d} atoms in total:") 

819 for region_id in np.unique(region): 

820 n_at = np.count_nonzero(region == region_id) 

821 print(f"{n_at:16d} {region_id}") 

822 

823 qm_atoms = atoms[self.qm_selection_mask] 

824 symbol_counts = qm_atoms.symbols.formula.count() 

825 print("QM atoms types:") 

826 for symbol, count in symbol_counts.items(): 

827 print(f"{count:16d} {symbol}") 

828 

829 return region 

830 

831 def set_masks_from_region(self, region): 

832 """ 

833 Sets masks from provided region array 

834 """ 

835 self.qm_selection_mask = region == "QM" 

836 buffer_mask = region == "buffer" 

837 

838 self.qm_buffer_mask = self.qm_selection_mask ^ buffer_mask 

839 

840 def export_extxyz(self, atoms=None, filename="qmmm_atoms.xyz"): 

841 """ 

842 exports the atoms to extended xyz file with additional "region" 

843 array keeping the mapping between QM, buffer and MM parts of 

844 the simulation 

845 """ 

846 if atoms is None: 

847 if self.atoms is None: 

848 raise ValueError('Calculator has no atoms') 

849 else: 

850 atoms = self.atoms 

851 

852 region = self.get_region_from_masks(atoms=atoms) 

853 

854 atoms_copy = atoms.copy() 

855 atoms_copy.new_array("region", region) 

856 

857 atoms_copy.calc = self # to keep the calculation results 

858 

859 atoms_copy.write(filename, format='extxyz') 

860 

861 @classmethod 

862 def import_extxyz(cls, filename, qm_calc, mm_calc): 

863 """ 

864 A static method to import the the mapping from an estxyz file saved by 

865 export_extxyz() function 

866 Parameters 

867 ---------- 

868 filename: string 

869 filename with saved configuration 

870 

871 qm_calc: Calculator object 

872 QM-calculator. 

873 mm_calc: Calculator object 

874 MM-calculator (should be scaled, see :class:`RescaledCalculator`) 

875 Can use `ForceConstantCalculator` based on QM force constants, if 

876 available. 

877 

878 Returns 

879 ------- 

880 New object of ForceQMMM calculator with qm_selection_mask and 

881 qm_buffer_mask set according to the region array in the saved file 

882 """ 

883 

884 from ase.io import read 

885 atoms = read(filename, format='extxyz') 

886 

887 if "region" in atoms.arrays: 

888 region = atoms.get_array("region") 

889 else: 

890 raise RuntimeError("Please provide extxyz file with 'region' array") 

891 

892 dummy_qm_mask = np.full_like(atoms, False, dtype=bool) 

893 dummy_qm_mask[0] = True 

894 

895 self = cls(atoms, dummy_qm_mask, 

896 qm_calc, mm_calc, buffer_width=1.0) 

897 

898 self.set_masks_from_region(region) 

899 

900 return self