Coverage for /builds/debichem-team/python-ase/ase/calculators/qmmm.py: 90.99%

455 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-03-06 04:00 +0000

1from typing import Sequence 

2 

3import numpy as np 

4 

5from ase.calculators.calculator import Calculator 

6from ase.cell import Cell 

7from ase.data import atomic_numbers 

8from ase.geometry import get_distances 

9from ase.parallel import world 

10from ase.utils import IOContext 

11 

12 

13class SimpleQMMM(Calculator): 

14 """Simple QMMM calculator.""" 

15 

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

17 

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

19 """SimpleQMMM object. 

20 

21 The energy is calculated as:: 

22 

23 _ _ _ 

24 E = E (R ) - E (R ) + E (R ) 

25 QM QM MM QM MM all 

26 

27 parameters: 

28 

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

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

31 qmcalc: Calculator object 

32 QM-calculator. 

33 mmcalc1: Calculator object 

34 MM-calculator used for QM region. 

35 mmcalc2: Calculator object 

36 MM-calculator used for everything. 

37 vacuum: float or None 

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

39 calculator doesn't need a box. 

40 

41 """ 

42 self.selection = selection 

43 self.qmcalc = qmcalc 

44 self.mmcalc1 = mmcalc1 

45 self.mmcalc2 = mmcalc2 

46 self.vacuum = vacuum 

47 

48 self.qmatoms = None 

49 self.center = None 

50 

51 Calculator.__init__(self) 

52 

53 def _get_name(self): 

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

55 

56 def initialize_qm(self, atoms): 

57 constraints = atoms.constraints 

58 atoms.constraints = [] 

59 self.qmatoms = atoms[self.selection] 

60 atoms.constraints = constraints 

61 self.qmatoms.pbc = False 

62 if self.vacuum: 

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

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

65 

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

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

68 

69 if self.qmatoms is None: 

70 self.initialize_qm(atoms) 

71 

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

73 if self.vacuum: 

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

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

76 

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

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

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

80 forces = self.mmcalc2.get_forces(atoms) 

81 

82 if self.vacuum: 

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

84 forces[self.selection] += qmforces 

85 

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

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

88 

89 self.results['energy'] = energy 

90 self.results['forces'] = forces 

91 

92 

93class EIQMMM(Calculator, IOContext): 

94 """Explicit interaction QMMM calculator.""" 

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

96 

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

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

99 """EIQMMM object. 

100 

101 The energy is calculated as:: 

102 

103 _ _ _ _ 

104 E = E (R ) + E (R ) + E (R , R ) 

105 QM QM MM MM I QM MM 

106 

107 parameters: 

108 

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

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

111 qmcalc: Calculator object 

112 QM-calculator. 

113 mmcalc: Calculator object 

114 MM-calculator. 

115 interaction: Interaction object 

116 Interaction between QM and MM regions. 

117 vacuum: float or None 

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

119 calculator doesn't need a box. 

120 embedding: Embedding object or None 

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

122 default one. 

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

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

125 

126 """ 

127 

128 self.selection = selection 

129 

130 self.qmcalc = qmcalc 

131 self.mmcalc = mmcalc 

132 self.interaction = interaction 

133 self.vacuum = vacuum 

134 self.embedding = embedding 

135 

136 self.qmatoms = None 

137 self.mmatoms = None 

138 self.mask = None 

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

140 

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

142 

143 Calculator.__init__(self) 

144 

145 def _get_name(self): 

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

147 

148 def initialize(self, atoms): 

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

150 self.mask[self.selection] = True 

151 

152 constraints = atoms.constraints 

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

154 self.qmatoms = atoms[self.mask] 

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

156 atoms.constraints = constraints 

157 

158 self.qmatoms.pbc = False 

159 

160 if self.vacuum: 

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

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

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

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

165 

166 self.qmatoms.calc = self.qmcalc 

167 self.mmatoms.calc = self.mmcalc 

168 

169 if self.embedding is None: 

170 self.embedding = Embedding() 

171 

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

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

174 

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

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

177 

178 if self.qmatoms is None: 

179 self.initialize(atoms) 

180 

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

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

183 

184 if self.vacuum: 

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

186 self.qmatoms.positions += shift 

187 else: 

188 shift = (0, 0, 0) 

189 

190 self.embedding.update(shift) 

191 

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

193 self.qmatoms, self.mmatoms, shift) 

194 

195 qmenergy = self.qmatoms.get_potential_energy() 

196 mmenergy = self.mmatoms.get_potential_energy() 

197 energy = ienergy + qmenergy + mmenergy 

198 

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

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

201 

202 qmforces = self.qmatoms.get_forces() 

203 mmforces = self.mmatoms.get_forces() 

204 

205 mmforces += self.embedding.get_mm_forces() 

206 

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

208 forces[self.mask] = qmforces + iqmforces 

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

210 

211 self.results['energy'] = energy 

212 self.results['forces'] = forces 

213 

214 

215def wrap(D, cell, pbc): 

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

217 for i, periodic in enumerate(pbc): 

218 if periodic: 

219 d = D[:, i] 

220 L = cell[i] 

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

222 

223 

224class Embedding: 

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

226 """Point-charge embedding.""" 

227 self.qmatoms = None 

228 self.mmatoms = None 

229 self.molecule_size = molecule_size 

230 self.virtual_molecule_size = None 

231 self.parameters = parameters 

232 

233 def __repr__(self): 

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

235 

236 def initialize(self, qmatoms, mmatoms): 

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

238 self.qmatoms = qmatoms 

239 self.mmatoms = mmatoms 

240 charges = mmatoms.calc.get_virtual_charges(mmatoms) 

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

242 self.virtual_molecule_size = (self.molecule_size * 

243 len(charges) // len(mmatoms)) 

244 

245 def update(self, shift): 

246 """Update point-charge positions.""" 

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

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

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

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

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

252 mask1 = self.mmatoms.calc.mask 

253 mask2 = ~mask1 

254 vmask1 = self.mmatoms.calc.virtual_mask 

255 vmask2 = ~vmask1 

256 apm1 = self.mmatoms.calc.apm1 

257 apm2 = self.mmatoms.calc.apm2 

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

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

260 pos = self.mmatoms.positions 

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

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

263 pos = (pos1, pos2) 

264 else: 

265 pos = (self.mmatoms.positions, ) 

266 apm1 = self.molecule_size 

267 apm2 = self.molecule_size 

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

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

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

271 spm1 = self.mmatoms.calc.sites_per_mol 

272 spm2 = self.mmatoms.calc.sites_per_mol 

273 else: 

274 spm1 = self.molecule_size 

275 spm2 = spm1 

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

277 mask2 = mask1 

278 

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

280 com_all = [] 

281 apm = (apm1, apm2) 

282 mask = (mask1, mask2) 

283 spm = (spm1, spm2) 

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

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

286 

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

288 # each molecule: 

289 distances = positions[:, 0] - qmcenter 

290 

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

292 offsets = distances - positions[:, 0] 

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

294 

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

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

297 # Need per atom for C-code: 

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

299 com_all.append(com_pv) 

300 

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

302 

303 positions = wrap_pos.copy() 

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

305 

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

307 com_pv = np.zeros_like(positions) 

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

309 com_pv[m] = com_all[ii] 

310 

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

312 if 'rc2' in self.parameters: 

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

314 else: 

315 self.pcpot.set_positions(positions) 

316 

317 def get_mm_forces(self): 

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

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

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

321 

322 

323def combine_lj_lorenz_berthelot(sigmaqm, sigmamm, 

324 epsilonqm, epsilonmm): 

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

326 sigma = [] 

327 epsilon = [] 

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

329 if isinstance(sigmamm, Sequence): 

330 numcalcs = len(sigmamm) 

331 else: 

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

333 sigmamm = (sigmamm, ) 

334 epsilonmm = (epsilonmm, ) 

335 for cc in range(numcalcs): 

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

337 epsilon_c = np.zeros_like(sigma_c) 

338 

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

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

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

342 sigma.append(sigma_c) 

343 epsilon.append(epsilon_c) 

344 

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

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

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

348 

349 return sigma, epsilon 

350 

351 

352class LJInteractionsGeneral: 

353 name = 'LJ-general' 

354 

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

356 qm_molecule_size, mm_molecule_size=3, 

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

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

359 

360 sigmaqm: array 

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

362 subsystem 

363 epsilonqm: array 

364 As sigmaqm, but for epsilon-paramaters 

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

366 A (no counterions): 

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

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

369 B (counterions): 

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

371 the length of counterions in the MM subsystem, and 

372 arr2 is the array from A. 

373 epsilonmm: array or tuple 

374 As sigmamm but for epsilon-parameters. 

375 qm_molecule_size: int 

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

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

378 in the QM subsystem) 

379 mm_molecule_size: int 

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

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

382 calculator) 

383 

384 """ 

385 self.sigmaqm = sigmaqm 

386 self.epsilonqm = epsilonqm 

387 self.sigmamm = sigmamm 

388 self.epsilonmm = epsilonmm 

389 self.qms = qm_molecule_size 

390 self.mms = mm_molecule_size 

391 self.rc = rc 

392 self.width = width 

393 self.combine_lj() 

394 

395 def combine_lj(self): 

396 self.sigma, self.epsilon = combine_lj_lorenz_berthelot( 

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

398 

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

400 epsilon = self.epsilon 

401 sigma = self.sigma 

402 

403 # loop over possible multiple mm calculators 

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

405 apm1 = self.mms 

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

407 mask2 = mask1 

408 apm = (apm1, ) 

409 sigma = (sigma, ) 

410 epsilon = (epsilon, ) 

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

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

413 mask1 = mmatoms.calc.mask 

414 mask2 = ~mask1 

415 apm1 = mmatoms.calc.apm1 

416 apm2 = mmatoms.calc.apm2 

417 apm = (apm1, apm2) 

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

419 epsilon = epsilon[0] 

420 

421 mask = (mask1, mask2) 

422 e_all = 0 

423 qmforces_all = np.zeros_like(qmatoms.positions) 

424 mmforces_all = np.zeros_like(mmatoms.positions) 

425 

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

427 # cases of no counter ions. 

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

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

430 qmforces = np.zeros_like(qmatoms.positions) 

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

432 energy = 0.0 

433 

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

435 

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

437 # cutoff from first atom of each mol 

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

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

440 d00 = d002**0.5 

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

442 x2 = d00 < self.rc 

443 x12 = np.logical_and(x1, x2) 

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

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

446 t[x2] = 1.0 

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

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

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

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

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

452 continue 

453 R = mmpositions - qmpos[qa, :] 

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

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

456 c12 = c6**2 

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

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

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

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

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

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

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

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

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

466 

467 e_all += energy 

468 qmforces_all += qmforces 

469 mmforces_all[m] += mmforces 

470 

471 return e_all, qmforces_all, mmforces_all 

472 

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

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

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

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

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

478 

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

480 # each molecule: 

481 distances = positions[:, 0] - qmcenter 

482 

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

484 offsets = distances - positions[:, 0] 

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

486 

487 return positions 

488 

489 

490class LJInteractions: 

491 name = 'LJ' 

492 

493 def __init__(self, parameters): 

494 """Lennard-Jones type explicit interaction. 

495 

496 parameters: dict 

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

498 for that pair. 

499 

500 Example: 

501 

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

503 

504 """ 

505 self.parameters = {} 

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

507 Z1 = atomic_numbers[symbol1] 

508 Z2 = atomic_numbers[symbol2] 

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

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

511 

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

513 qmforces = np.zeros_like(qmatoms.positions) 

514 mmforces = np.zeros_like(mmatoms.positions) 

515 species = set(mmatoms.numbers) 

516 energy = 0.0 

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

518 for Z2 in species: 

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

520 continue 

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

522 mask = (mmatoms.numbers == Z2) 

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

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

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

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

527 c12 = c6**2 

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

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

530 F1 -= f.sum(0) 

531 mmforces[mask] += f 

532 return energy, qmforces, mmforces 

533 

534 

535class RescaledCalculator(Calculator): 

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

537 lattice constant and bulk modulus 

538 

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

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

541 for a derivation of the scaling constants. 

542 """ 

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

544 

545 def __init__(self, mm_calc, 

546 qm_lattice_constant, qm_bulk_modulus, 

547 mm_lattice_constant, mm_bulk_modulus): 

548 Calculator.__init__(self) 

549 self.mm_calc = mm_calc 

550 self.alpha = qm_lattice_constant / mm_lattice_constant 

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

552 

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

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

555 

556 # mm_pos = atoms.get_positions() 

557 scaled_atoms = atoms.copy() 

558 

559 # scaled_atoms.positions = mm_pos/self.alpha 

560 mm_cell = atoms.get_cell() 

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

562 

563 results = {} 

564 

565 if 'energy' in properties: 

566 energy = self.mm_calc.get_potential_energy(scaled_atoms) 

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

568 

569 if 'forces' in properties: 

570 forces = self.mm_calc.get_forces(scaled_atoms) 

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

572 

573 if 'stress' in properties: 

574 stress = self.mm_calc.get_stress(scaled_atoms) 

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

576 

577 self.results = results 

578 

579 

580class ForceConstantCalculator(Calculator): 

581 """ 

582 Compute forces based on provided force-constant matrix 

583 

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

585 of QM method. 

586 """ 

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

588 

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

590 """ 

591 Parameters: 

592 

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

594 Force constant matrix. 

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

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

597 ref: ase.atoms.Atoms 

598 Atoms object for reference configuration 

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

600 Value of forces at reference configuration 

601 """ 

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

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

604 self.D = D 

605 self.ref = ref 

606 self.f0 = f0 

607 self.size = len(ref) 

608 Calculator.__init__(self) 

609 

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

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

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

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

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

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

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

617 self.results['energy'] = 0.0 

618 

619 

620class ForceQMMM(Calculator): 

621 """ 

622 Force-based QM/MM calculator 

623 

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

625 with MM forces: 

626 

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

628 { F^i_MM otherwise 

629 

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

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

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

633 """ 

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

635 

636 def __init__(self, 

637 atoms, 

638 qm_selection_mask, 

639 qm_calc, 

640 mm_calc, 

641 buffer_width, 

642 vacuum=5., 

643 zero_mean=True, 

644 qm_cell_round_off=3, 

645 qm_radius=None): 

646 """ 

647 ForceQMMM calculator 

648 

649 Parameters: 

650 

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

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

653 qm_calc: Calculator object 

654 QM-calculator. 

655 mm_calc: Calculator object 

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

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

658 available. 

659 vacuum: float or None 

660 Amount of vacuum to add around QM atoms. 

661 zero_mean: bool 

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

663 qm_cell_round_off: float 

664 Tolerance value in Angstrom to round the qm cluster cell 

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

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

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

668 in qm region. 

669 """ 

670 

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

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

673 

674 self.qm_selection_mask = qm_selection_mask 

675 self.qm_calc = qm_calc 

676 self.mm_calc = mm_calc 

677 self.vacuum = vacuum 

678 self.buffer_width = buffer_width 

679 self.zero_mean = zero_mean 

680 self.qm_cell_round_off = qm_cell_round_off 

681 self.qm_radius = qm_radius 

682 

683 self.qm_buffer_mask = None 

684 

685 Calculator.__init__(self) 

686 

687 def initialize_qm_buffer_mask(self, atoms): 

688 """ 

689 Initialises system to perform qm calculation 

690 """ 

691 # calculate the distances between all atoms and qm atoms 

692 # qm_distance_matrix is a [N_QM_atoms x N_atoms] matrix 

693 _, qm_distance_matrix = get_distances( 

694 atoms.positions[self.qm_selection_mask], 

695 atoms.positions, 

696 atoms.cell, atoms.pbc) 

697 

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

699 

700 # every r_qm is a matrix of distances 

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

702 for r_qm in qm_distance_matrix: 

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

704 

705 def get_qm_cluster(self, atoms): 

706 

707 if self.qm_buffer_mask is None: 

708 self.initialize_qm_buffer_mask(atoms) 

709 

710 qm_cluster = atoms[self.qm_buffer_mask] 

711 del qm_cluster.constraints 

712 

713 round_cell = False 

714 if self.qm_radius is None: 

715 round_cell = True 

716 # get all distances between qm atoms. 

717 # Treat all X, Y and Z directions independently 

718 # only distance between qm atoms is calculated 

719 # in order to estimate qm radius in thee directions 

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

721 cell=atoms.cell, pbc=atoms.pbc) 

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

723 # of max distance between qm atoms 

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

725 

726 if atoms.cell.orthorhombic: 

727 cell_size = np.diagonal(atoms.cell) 

728 else: 

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

730 

731 # check if qm_cluster should be left periodic 

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

733 # otherwise change to non pbc 

734 # and make a cluster in a vacuum configuration 

735 qm_cluster_pbc = (atoms.pbc & 

736 (cell_size < 

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

738 

739 # start with the original orthorhombic cell 

740 qm_cluster_cell = cell_size.copy() 

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

742 qm_cluster_cell[~qm_cluster_pbc] = ( 

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

744 self.buffer_width + 

745 self.vacuum)) 

746 

747 if round_cell: 

748 # round the qm cell to the required tolerance 

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

750 (qm_cluster_cell[~qm_cluster_pbc]) / 

751 self.qm_cell_round_off) * 

752 self.qm_cell_round_off) 

753 

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

755 qm_cluster.pbc = qm_cluster_pbc 

756 

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

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

759 

760 if 'cell_origin' in qm_cluster.info: 

761 del qm_cluster.info['cell_origin'] 

762 

763 # center the cluster only in non pbc directions 

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

765 

766 return qm_cluster 

767 

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

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

770 

771 qm_cluster = self.get_qm_cluster(atoms) 

772 

773 forces = self.mm_calc.get_forces(atoms) 

774 qm_forces = self.qm_calc.get_forces(qm_cluster) 

775 

776 forces[self.qm_selection_mask] = \ 

777 qm_forces[self.qm_selection_mask[self.qm_buffer_mask]] 

778 

779 if self.zero_mean: 

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

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

782 

783 self.results['forces'] = forces 

784 self.results['energy'] = 0.0 

785 

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

787 """ 

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

789 the array are: 

790 QM - qm atoms 

791 buffer - buffer atoms 

792 MM - atoms treated with mm calculator 

793 """ 

794 if atoms is None: 

795 if self.atoms is None: 

796 raise ValueError('Calculator has no atoms') 

797 else: 

798 atoms = self.atoms 

799 

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

801 

802 region[self.qm_selection_mask] = ( 

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

804 

805 buffer_only_mask = self.qm_buffer_mask & ~self.qm_selection_mask 

806 

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

808 "buffer") 

809 

810 if print_mapping: 

811 

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

813 for region_id in np.unique(region): 

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

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

816 

817 qm_atoms = atoms[self.qm_selection_mask] 

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

819 print("QM atoms types:") 

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

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

822 

823 return region 

824 

825 def set_masks_from_region(self, region): 

826 """ 

827 Sets masks from provided region array 

828 """ 

829 self.qm_selection_mask = region == "QM" 

830 buffer_mask = region == "buffer" 

831 

832 self.qm_buffer_mask = self.qm_selection_mask ^ buffer_mask 

833 

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

835 """ 

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

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

838 the simulation 

839 """ 

840 if atoms is None: 

841 if self.atoms is None: 

842 raise ValueError('Calculator has no atoms') 

843 else: 

844 atoms = self.atoms 

845 

846 region = self.get_region_from_masks(atoms=atoms) 

847 

848 atoms_copy = atoms.copy() 

849 atoms_copy.new_array("region", region) 

850 

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

852 

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

854 

855 @classmethod 

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

857 """ 

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

859 export_extxyz() function 

860 Parameters 

861 ---------- 

862 filename: string 

863 filename with saved configuration 

864 

865 qm_calc: Calculator object 

866 QM-calculator. 

867 mm_calc: Calculator object 

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

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

870 available. 

871 

872 Returns 

873 ------- 

874 New object of ForceQMMM calculator with qm_selection_mask and 

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

876 """ 

877 

878 from ase.io import read 

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

880 

881 if "region" in atoms.arrays: 

882 region = atoms.get_array("region") 

883 else: 

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

885 

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

887 dummy_qm_mask[0] = True 

888 

889 self = cls(atoms, dummy_qm_mask, 

890 qm_calc, mm_calc, buffer_width=1.0) 

891 

892 self.set_masks_from_region(region) 

893 

894 return self