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

455 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-08-02 00:12 +0000

1# fmt: off 

2 

3from typing 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 selection: list of int, slice object or list of bool 

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

33 qmcalc: Calculator object 

34 QM-calculator. 

35 mmcalc1: Calculator object 

36 MM-calculator used for QM region. 

37 mmcalc2: Calculator object 

38 MM-calculator used for everything. 

39 vacuum: float or None 

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

41 calculator doesn't need a box. 

42 

43 """ 

44 self.selection = selection 

45 self.qmcalc = qmcalc 

46 self.mmcalc1 = mmcalc1 

47 self.mmcalc2 = mmcalc2 

48 self.vacuum = vacuum 

49 

50 self.qmatoms = None 

51 self.center = None 

52 

53 Calculator.__init__(self) 

54 

55 def _get_name(self): 

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

57 

58 def initialize_qm(self, atoms): 

59 constraints = atoms.constraints 

60 atoms.constraints = [] 

61 self.qmatoms = atoms[self.selection] 

62 atoms.constraints = constraints 

63 self.qmatoms.pbc = False 

64 if self.vacuum: 

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

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

67 

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

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

70 

71 if self.qmatoms is None: 

72 self.initialize_qm(atoms) 

73 

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

75 if self.vacuum: 

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

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

78 

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

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

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

82 forces = self.mmcalc2.get_forces(atoms) 

83 

84 if self.vacuum: 

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

86 forces[self.selection] += qmforces 

87 

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

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

90 

91 self.results['energy'] = energy 

92 self.results['forces'] = forces 

93 

94 

95class EIQMMM(Calculator, IOContext): 

96 """Explicit interaction QMMM calculator.""" 

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

98 

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

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

101 """EIQMMM object. 

102 

103 The energy is calculated as:: 

104 

105 _ _ _ _ 

106 E = E (R ) + E (R ) + E (R , R ) 

107 QM QM MM MM I QM MM 

108 

109 parameters: 

110 

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

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

113 qmcalc: Calculator object 

114 QM-calculator. 

115 mmcalc: Calculator object 

116 MM-calculator. 

117 interaction: Interaction object 

118 Interaction between QM and MM regions. 

119 vacuum: float or None 

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

121 calculator doesn't need a box. 

122 embedding: Embedding object or None 

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

124 default one. 

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

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

127 

128 """ 

129 

130 self.selection = selection 

131 

132 self.qmcalc = qmcalc 

133 self.mmcalc = mmcalc 

134 self.interaction = interaction 

135 self.vacuum = vacuum 

136 self.embedding = embedding 

137 

138 self.qmatoms = None 

139 self.mmatoms = None 

140 self.mask = None 

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

142 

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

144 

145 Calculator.__init__(self) 

146 

147 def _get_name(self): 

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

149 

150 def initialize(self, atoms): 

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

152 self.mask[self.selection] = True 

153 

154 constraints = atoms.constraints 

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

156 self.qmatoms = atoms[self.mask] 

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

158 atoms.constraints = constraints 

159 

160 self.qmatoms.pbc = False 

161 

162 if self.vacuum: 

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

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

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

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

167 

168 self.qmatoms.calc = self.qmcalc 

169 self.mmatoms.calc = self.mmcalc 

170 

171 if self.embedding is None: 

172 self.embedding = Embedding() 

173 

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

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

176 

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

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

179 

180 if self.qmatoms is None: 

181 self.initialize(atoms) 

182 

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

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

185 

186 if self.vacuum: 

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

188 self.qmatoms.positions += shift 

189 else: 

190 shift = (0, 0, 0) 

191 

192 self.embedding.update(shift) 

193 

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

195 self.qmatoms, self.mmatoms, shift) 

196 

197 qmenergy = self.qmatoms.get_potential_energy() 

198 mmenergy = self.mmatoms.get_potential_energy() 

199 energy = ienergy + qmenergy + mmenergy 

200 

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

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

203 

204 qmforces = self.qmatoms.get_forces() 

205 mmforces = self.mmatoms.get_forces() 

206 

207 mmforces += self.embedding.get_mm_forces() 

208 

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

210 forces[self.mask] = qmforces + iqmforces 

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

212 

213 self.results['energy'] = energy 

214 self.results['forces'] = forces 

215 

216 

217def wrap(D, cell, pbc): 

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

219 for i, periodic in enumerate(pbc): 

220 if periodic: 

221 d = D[:, i] 

222 L = cell[i] 

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

224 

225 

226class Embedding: 

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

228 """Point-charge embedding.""" 

229 self.qmatoms = None 

230 self.mmatoms = None 

231 self.molecule_size = molecule_size 

232 self.virtual_molecule_size = None 

233 self.parameters = parameters 

234 

235 def __repr__(self): 

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

237 

238 def initialize(self, qmatoms, mmatoms): 

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

240 self.qmatoms = qmatoms 

241 self.mmatoms = mmatoms 

242 charges = mmatoms.calc.get_virtual_charges(mmatoms) 

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

244 self.virtual_molecule_size = (self.molecule_size * 

245 len(charges) // len(mmatoms)) 

246 

247 def update(self, shift): 

248 """Update point-charge positions.""" 

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

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

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

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

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

254 mask1 = self.mmatoms.calc.mask 

255 mask2 = ~mask1 

256 vmask1 = self.mmatoms.calc.virtual_mask 

257 vmask2 = ~vmask1 

258 apm1 = self.mmatoms.calc.apm1 

259 apm2 = self.mmatoms.calc.apm2 

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

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

262 pos = self.mmatoms.positions 

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

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

265 pos = (pos1, pos2) 

266 else: 

267 pos = (self.mmatoms.positions, ) 

268 apm1 = self.molecule_size 

269 apm2 = self.molecule_size 

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

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

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

273 spm1 = self.mmatoms.calc.sites_per_mol 

274 spm2 = self.mmatoms.calc.sites_per_mol 

275 else: 

276 spm1 = self.molecule_size 

277 spm2 = spm1 

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

279 mask2 = mask1 

280 

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

282 com_all = [] 

283 apm = (apm1, apm2) 

284 mask = (mask1, mask2) 

285 spm = (spm1, spm2) 

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

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

288 

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

290 # each molecule: 

291 distances = positions[:, 0] - qmcenter 

292 

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

294 offsets = distances - positions[:, 0] 

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

296 

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

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

299 # Need per atom for C-code: 

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

301 com_all.append(com_pv) 

302 

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

304 

305 positions = wrap_pos.copy() 

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

307 

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

309 com_pv = np.zeros_like(positions) 

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

311 com_pv[m] = com_all[ii] 

312 

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

314 if 'rc2' in self.parameters: 

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

316 else: 

317 self.pcpot.set_positions(positions) 

318 

319 def get_mm_forces(self): 

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

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

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

323 

324 

325def combine_lj_lorenz_berthelot(sigmaqm, sigmamm, 

326 epsilonqm, epsilonmm): 

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

328 sigma = [] 

329 epsilon = [] 

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

331 if isinstance(sigmamm, Sequence): 

332 numcalcs = len(sigmamm) 

333 else: 

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

335 sigmamm = (sigmamm, ) 

336 epsilonmm = (epsilonmm, ) 

337 for cc in range(numcalcs): 

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

339 epsilon_c = np.zeros_like(sigma_c) 

340 

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

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

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

344 sigma.append(sigma_c) 

345 epsilon.append(epsilon_c) 

346 

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

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

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

350 

351 return sigma, epsilon 

352 

353 

354class LJInteractionsGeneral: 

355 name = 'LJ-general' 

356 

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

358 qm_molecule_size, mm_molecule_size=3, 

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

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

361 

362 sigmaqm: array 

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

364 subsystem 

365 epsilonqm: array 

366 As sigmaqm, but for epsilon-paramaters 

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

368 A (no counterions): 

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

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

371 B (counterions): 

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

373 the length of counterions in the MM subsystem, and 

374 arr2 is the array from A. 

375 epsilonmm: array or tuple 

376 As sigmamm but for epsilon-parameters. 

377 qm_molecule_size: int 

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

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

380 in the QM subsystem) 

381 mm_molecule_size: int 

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

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

384 calculator) 

385 

386 """ 

387 self.sigmaqm = sigmaqm 

388 self.epsilonqm = epsilonqm 

389 self.sigmamm = sigmamm 

390 self.epsilonmm = epsilonmm 

391 self.qms = qm_molecule_size 

392 self.mms = mm_molecule_size 

393 self.rc = rc 

394 self.width = width 

395 self.combine_lj() 

396 

397 def combine_lj(self): 

398 self.sigma, self.epsilon = combine_lj_lorenz_berthelot( 

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

400 

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

402 epsilon = self.epsilon 

403 sigma = self.sigma 

404 

405 # loop over possible multiple mm calculators 

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

407 apm1 = self.mms 

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

409 mask2 = mask1 

410 apm = (apm1, ) 

411 sigma = (sigma, ) 

412 epsilon = (epsilon, ) 

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

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

415 mask1 = mmatoms.calc.mask 

416 mask2 = ~mask1 

417 apm1 = mmatoms.calc.apm1 

418 apm2 = mmatoms.calc.apm2 

419 apm = (apm1, apm2) 

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

421 epsilon = epsilon[0] 

422 

423 mask = (mask1, mask2) 

424 e_all = 0 

425 qmforces_all = np.zeros_like(qmatoms.positions) 

426 mmforces_all = np.zeros_like(mmatoms.positions) 

427 

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

429 # cases of no counter ions. 

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

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

432 qmforces = np.zeros_like(qmatoms.positions) 

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

434 energy = 0.0 

435 

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

437 

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

439 # cutoff from first atom of each mol 

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

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

442 d00 = d002**0.5 

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

444 x2 = d00 < self.rc 

445 x12 = np.logical_and(x1, x2) 

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

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

448 t[x2] = 1.0 

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

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

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

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

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

454 continue 

455 R = mmpositions - qmpos[qa, :] 

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

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

458 c12 = c6**2 

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

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

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

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

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

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

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

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

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

468 

469 e_all += energy 

470 qmforces_all += qmforces 

471 mmforces_all[m] += mmforces 

472 

473 return e_all, qmforces_all, mmforces_all 

474 

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

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

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

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

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

480 

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

482 # each molecule: 

483 distances = positions[:, 0] - qmcenter 

484 

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

486 offsets = distances - positions[:, 0] 

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

488 

489 return positions 

490 

491 

492class LJInteractions: 

493 name = 'LJ' 

494 

495 def __init__(self, parameters): 

496 """Lennard-Jones type explicit interaction. 

497 

498 parameters: dict 

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

500 for that pair. 

501 

502 Example: 

503 

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

505 

506 """ 

507 self.parameters = {} 

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

509 Z1 = atomic_numbers[symbol1] 

510 Z2 = atomic_numbers[symbol2] 

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

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

513 

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

515 qmforces = np.zeros_like(qmatoms.positions) 

516 mmforces = np.zeros_like(mmatoms.positions) 

517 species = set(mmatoms.numbers) 

518 energy = 0.0 

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

520 for Z2 in species: 

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

522 continue 

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

524 mask = (mmatoms.numbers == Z2) 

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

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

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

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

529 c12 = c6**2 

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

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

532 F1 -= f.sum(0) 

533 mmforces[mask] += f 

534 return energy, qmforces, mmforces 

535 

536 

537class RescaledCalculator(Calculator): 

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

539 lattice constant and bulk modulus 

540 

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

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

543 for a derivation of the scaling constants. 

544 """ 

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

546 

547 def __init__(self, mm_calc, 

548 qm_lattice_constant, qm_bulk_modulus, 

549 mm_lattice_constant, mm_bulk_modulus): 

550 Calculator.__init__(self) 

551 self.mm_calc = mm_calc 

552 self.alpha = qm_lattice_constant / mm_lattice_constant 

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

554 

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

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

557 

558 # mm_pos = atoms.get_positions() 

559 scaled_atoms = atoms.copy() 

560 

561 # scaled_atoms.positions = mm_pos/self.alpha 

562 mm_cell = atoms.get_cell() 

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

564 

565 results = {} 

566 

567 if 'energy' in properties: 

568 energy = self.mm_calc.get_potential_energy(scaled_atoms) 

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

570 

571 if 'forces' in properties: 

572 forces = self.mm_calc.get_forces(scaled_atoms) 

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

574 

575 if 'stress' in properties: 

576 stress = self.mm_calc.get_stress(scaled_atoms) 

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

578 

579 self.results = results 

580 

581 

582class ForceConstantCalculator(Calculator): 

583 """ 

584 Compute forces based on provided force-constant matrix 

585 

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

587 of QM method. 

588 """ 

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

590 

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

592 """ 

593 Parameters: 

594 

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

596 Force constant matrix. 

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

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

599 ref: ase.atoms.Atoms 

600 Atoms object for reference configuration 

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

602 Value of forces at reference configuration 

603 """ 

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

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

606 self.D = D 

607 self.ref = ref 

608 self.f0 = f0 

609 self.size = len(ref) 

610 Calculator.__init__(self) 

611 

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

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

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

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

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

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

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

619 self.results['energy'] = 0.0 

620 

621 

622class ForceQMMM(Calculator): 

623 """ 

624 Force-based QM/MM calculator 

625 

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

627 with MM forces: 

628 

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

630 { F^i_MM otherwise 

631 

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

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

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

635 """ 

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

637 

638 def __init__(self, 

639 atoms, 

640 qm_selection_mask, 

641 qm_calc, 

642 mm_calc, 

643 buffer_width, 

644 vacuum=5., 

645 zero_mean=True, 

646 qm_cell_round_off=3, 

647 qm_radius=None): 

648 """ 

649 ForceQMMM calculator 

650 

651 Parameters: 

652 

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

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

655 qm_calc: Calculator object 

656 QM-calculator. 

657 mm_calc: Calculator object 

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

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

660 available. 

661 vacuum: float or None 

662 Amount of vacuum to add around QM atoms. 

663 zero_mean: bool 

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

665 qm_cell_round_off: float 

666 Tolerance value in Angstrom to round the qm cluster cell 

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

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

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

670 in qm region. 

671 """ 

672 

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

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

675 

676 self.qm_selection_mask = qm_selection_mask 

677 self.qm_calc = qm_calc 

678 self.mm_calc = mm_calc 

679 self.vacuum = vacuum 

680 self.buffer_width = buffer_width 

681 self.zero_mean = zero_mean 

682 self.qm_cell_round_off = qm_cell_round_off 

683 self.qm_radius = qm_radius 

684 

685 self.qm_buffer_mask = None 

686 

687 Calculator.__init__(self) 

688 

689 def initialize_qm_buffer_mask(self, atoms): 

690 """ 

691 Initialises system to perform qm calculation 

692 """ 

693 # calculate the distances between all atoms and qm atoms 

694 # qm_distance_matrix is a [N_QM_atoms x N_atoms] matrix 

695 _, qm_distance_matrix = get_distances( 

696 atoms.positions[self.qm_selection_mask], 

697 atoms.positions, 

698 atoms.cell, atoms.pbc) 

699 

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

701 

702 # every r_qm is a matrix of distances 

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

704 for r_qm in qm_distance_matrix: 

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

706 

707 def get_qm_cluster(self, atoms): 

708 

709 if self.qm_buffer_mask is None: 

710 self.initialize_qm_buffer_mask(atoms) 

711 

712 qm_cluster = atoms[self.qm_buffer_mask] 

713 del qm_cluster.constraints 

714 

715 round_cell = False 

716 if self.qm_radius is None: 

717 round_cell = True 

718 # get all distances between qm atoms. 

719 # Treat all X, Y and Z directions independently 

720 # only distance between qm atoms is calculated 

721 # in order to estimate qm radius in thee directions 

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

723 cell=atoms.cell, pbc=atoms.pbc) 

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

725 # of max distance between qm atoms 

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

727 

728 if atoms.cell.orthorhombic: 

729 cell_size = np.diagonal(atoms.cell) 

730 else: 

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

732 

733 # check if qm_cluster should be left periodic 

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

735 # otherwise change to non pbc 

736 # and make a cluster in a vacuum configuration 

737 qm_cluster_pbc = (atoms.pbc & 

738 (cell_size < 

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

740 

741 # start with the original orthorhombic cell 

742 qm_cluster_cell = cell_size.copy() 

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

744 qm_cluster_cell[~qm_cluster_pbc] = ( 

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

746 self.buffer_width + 

747 self.vacuum)) 

748 

749 if round_cell: 

750 # round the qm cell to the required tolerance 

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

752 (qm_cluster_cell[~qm_cluster_pbc]) / 

753 self.qm_cell_round_off) * 

754 self.qm_cell_round_off) 

755 

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

757 qm_cluster.pbc = qm_cluster_pbc 

758 

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

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

761 

762 if 'cell_origin' in qm_cluster.info: 

763 del qm_cluster.info['cell_origin'] 

764 

765 # center the cluster only in non pbc directions 

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

767 

768 return qm_cluster 

769 

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

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

772 

773 qm_cluster = self.get_qm_cluster(atoms) 

774 

775 forces = self.mm_calc.get_forces(atoms) 

776 qm_forces = self.qm_calc.get_forces(qm_cluster) 

777 

778 forces[self.qm_selection_mask] = \ 

779 qm_forces[self.qm_selection_mask[self.qm_buffer_mask]] 

780 

781 if self.zero_mean: 

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

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

784 

785 self.results['forces'] = forces 

786 self.results['energy'] = 0.0 

787 

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

789 """ 

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

791 the array are: 

792 QM - qm atoms 

793 buffer - buffer atoms 

794 MM - atoms treated with mm calculator 

795 """ 

796 if atoms is None: 

797 if self.atoms is None: 

798 raise ValueError('Calculator has no atoms') 

799 else: 

800 atoms = self.atoms 

801 

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

803 

804 region[self.qm_selection_mask] = ( 

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

806 

807 buffer_only_mask = self.qm_buffer_mask & ~self.qm_selection_mask 

808 

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

810 "buffer") 

811 

812 if print_mapping: 

813 

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

815 for region_id in np.unique(region): 

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

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

818 

819 qm_atoms = atoms[self.qm_selection_mask] 

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

821 print("QM atoms types:") 

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

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

824 

825 return region 

826 

827 def set_masks_from_region(self, region): 

828 """ 

829 Sets masks from provided region array 

830 """ 

831 self.qm_selection_mask = region == "QM" 

832 buffer_mask = region == "buffer" 

833 

834 self.qm_buffer_mask = self.qm_selection_mask ^ buffer_mask 

835 

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

837 """ 

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

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

840 the simulation 

841 """ 

842 if atoms is None: 

843 if self.atoms is None: 

844 raise ValueError('Calculator has no atoms') 

845 else: 

846 atoms = self.atoms 

847 

848 region = self.get_region_from_masks(atoms=atoms) 

849 

850 atoms_copy = atoms.copy() 

851 atoms_copy.new_array("region", region) 

852 

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

854 

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

856 

857 @classmethod 

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

859 """ 

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

861 export_extxyz() function 

862 Parameters 

863 ---------- 

864 filename: string 

865 filename with saved configuration 

866 

867 qm_calc: Calculator object 

868 QM-calculator. 

869 mm_calc: Calculator object 

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

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

872 available. 

873 

874 Returns 

875 ------- 

876 New object of ForceQMMM calculator with qm_selection_mask and 

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

878 """ 

879 

880 from ase.io import read 

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

882 

883 if "region" in atoms.arrays: 

884 region = atoms.get_array("region") 

885 else: 

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

887 

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

889 dummy_qm_mask[0] = True 

890 

891 self = cls(atoms, dummy_qm_mask, 

892 qm_calc, mm_calc, buffer_width=1.0) 

893 

894 self.set_masks_from_region(region) 

895 

896 return self