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
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +0000
1from typing import Sequence
3import numpy as np
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
13class SimpleQMMM(Calculator):
14 """Simple QMMM calculator."""
16 implemented_properties = ['energy', 'forces']
18 def __init__(self, selection, qmcalc, mmcalc1, mmcalc2, vacuum=None):
19 """SimpleQMMM object.
21 The energy is calculated as::
23 _ _ _
24 E = E (R ) - E (R ) + E (R )
25 QM QM MM QM MM all
27 parameters:
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.
41 """
42 self.selection = selection
43 self.qmcalc = qmcalc
44 self.mmcalc1 = mmcalc1
45 self.mmcalc2 = mmcalc2
46 self.vacuum = vacuum
48 self.qmatoms = None
49 self.center = None
51 Calculator.__init__(self)
53 def _get_name(self):
54 return f'{self.qmcalc.name}-{self.mmcalc1.name}+{self.mmcalc1.name}'
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)
66 def calculate(self, atoms, properties, system_changes):
67 Calculator.calculate(self, atoms, properties, system_changes)
69 if self.qmatoms is None:
70 self.initialize_qm(atoms)
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))
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)
82 if self.vacuum:
83 qmforces -= qmforces.mean(axis=0)
84 forces[self.selection] += qmforces
86 energy -= self.mmcalc1.get_potential_energy(self.qmatoms)
87 forces[self.selection] -= self.mmcalc1.get_forces(self.qmatoms)
89 self.results['energy'] = energy
90 self.results['forces'] = forces
93class EIQMMM(Calculator, IOContext):
94 """Explicit interaction QMMM calculator."""
95 implemented_properties = ['energy', 'forces']
97 def __init__(self, selection, qmcalc, mmcalc, interaction,
98 vacuum=None, embedding=None, output=None, comm=world):
99 """EIQMMM object.
101 The energy is calculated as::
103 _ _ _ _
104 E = E (R ) + E (R ) + E (R , R )
105 QM QM MM MM I QM MM
107 parameters:
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).
126 """
128 self.selection = selection
130 self.qmcalc = qmcalc
131 self.mmcalc = mmcalc
132 self.interaction = interaction
133 self.vacuum = vacuum
134 self.embedding = embedding
136 self.qmatoms = None
137 self.mmatoms = None
138 self.mask = None
139 self.center = None # center of QM atoms in QM-box
141 self.output = self.openfile(file=output, comm=comm)
143 Calculator.__init__(self)
145 def _get_name(self):
146 return f'{self.qmcalc.name}+{self.interaction.name}+{self.mmcalc.name}'
148 def initialize(self, atoms):
149 self.mask = np.zeros(len(atoms), bool)
150 self.mask[self.selection] = True
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
158 self.qmatoms.pbc = False
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)
166 self.qmatoms.calc = self.qmcalc
167 self.mmatoms.calc = self.mmcalc
169 if self.embedding is None:
170 self.embedding = Embedding()
172 self.embedding.initialize(self.qmatoms, self.mmatoms)
173 print('Embedding:', self.embedding, file=self.output)
175 def calculate(self, atoms, properties, system_changes):
176 Calculator.calculate(self, atoms, properties, system_changes)
178 if self.qmatoms is None:
179 self.initialize(atoms)
181 self.mmatoms.set_positions(atoms.positions[~self.mask])
182 self.qmatoms.set_positions(atoms.positions[self.mask])
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)
190 self.embedding.update(shift)
192 ienergy, iqmforces, immforces = self.interaction.calculate(
193 self.qmatoms, self.mmatoms, shift)
195 qmenergy = self.qmatoms.get_potential_energy()
196 mmenergy = self.mmatoms.get_potential_energy()
197 energy = ienergy + qmenergy + mmenergy
199 print('Energies: {:12.3f} {:+12.3f} {:+12.3f} = {:12.3f}'
200 .format(ienergy, qmenergy, mmenergy, energy), file=self.output)
202 qmforces = self.qmatoms.get_forces()
203 mmforces = self.mmatoms.get_forces()
205 mmforces += self.embedding.get_mm_forces()
207 forces = np.empty((len(atoms), 3))
208 forces[self.mask] = qmforces + iqmforces
209 forces[~self.mask] = mmforces + immforces
211 self.results['energy'] = energy
212 self.results['forces'] = forces
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
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
233 def __repr__(self):
234 return f'Embedding(molecule_size={self.molecule_size})'
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))
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
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
287 # Distances from the center of the QM box to the first atom of
288 # each molecule:
289 distances = positions[:, 0] - qmcenter
291 wrap(distances, self.mmatoms.cell.diagonal(), self.mmatoms.pbc)
292 offsets = distances - positions[:, 0]
293 positions += offsets[:, np.newaxis] + qmcenter
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)
301 wrap_pos[m] = positions.reshape((-1, 3))
303 positions = wrap_pos.copy()
304 positions = self.mmatoms.calc.add_virtual_sites(positions)
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]
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)
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)
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)
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)
345 if numcalcs == 1: # retain original, 1 calc function
346 sigma = np.array(sigma[0])
347 epsilon = np.array(epsilon[0])
349 return sigma, epsilon
352class LJInteractionsGeneral:
353 name = 'LJ-general'
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.
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)
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()
395 def combine_lj(self):
396 self.sigma, self.epsilon = combine_lj_lorenz_berthelot(
397 self.sigmaqm, self.sigmamm, self.epsilonqm, self.epsilonmm)
399 def calculate(self, qmatoms, mmatoms, shift):
400 epsilon = self.epsilon
401 sigma = self.sigma
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]
421 mask = (mask1, mask2)
422 e_all = 0
423 qmforces_all = np.zeros_like(qmatoms.positions)
424 mmforces_all = np.zeros_like(mmatoms.positions)
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
434 qmpositions = qmatoms.positions.reshape((-1, self.qms, 3))
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
467 e_all += energy
468 qmforces_all += qmforces
469 mmforces_all[m] += mmforces
471 return e_all, qmforces_all, mmforces_all
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
479 # Distances from the center of the QM box to the first atom of
480 # each molecule:
481 distances = positions[:, 0] - qmcenter
483 wrap(distances, mmatoms.cell.diagonal(), mmatoms.pbc)
484 offsets = distances - positions[:, 0]
485 positions += offsets[:, np.newaxis] + qmcenter
487 return positions
490class LJInteractions:
491 name = 'LJ'
493 def __init__(self, parameters):
494 """Lennard-Jones type explicit interaction.
496 parameters: dict
497 Mapping from pair of atoms to tuple containing epsilon and sigma
498 for that pair.
500 Example:
502 lj = LJInteractions({('O', 'O'): (eps, sigma)})
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
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
535class RescaledCalculator(Calculator):
536 """Rescales length and energy of a calculators to match given
537 lattice constant and bulk modulus
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']
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)
553 def calculate(self, atoms, properties, system_changes):
554 Calculator.calculate(self, atoms, properties, system_changes)
556 # mm_pos = atoms.get_positions()
557 scaled_atoms = atoms.copy()
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)
563 results = {}
565 if 'energy' in properties:
566 energy = self.mm_calc.get_potential_energy(scaled_atoms)
567 results['energy'] = energy / self.beta
569 if 'forces' in properties:
570 forces = self.mm_calc.get_forces(scaled_atoms)
571 results['forces'] = forces / (self.beta * self.alpha)
573 if 'stress' in properties:
574 stress = self.mm_calc.get_stress(scaled_atoms)
575 results['stress'] = stress / (self.beta * self.alpha**3)
577 self.results = results
580class ForceConstantCalculator(Calculator):
581 """
582 Compute forces based on provided force-constant matrix
584 Useful with `ForceQMMM` to do harmonic QM/MM using force constants
585 of QM method.
586 """
587 implemented_properties = ['forces', 'energy']
589 def __init__(self, D, ref, f0):
590 """
591 Parameters:
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)
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
620class ForceQMMM(Calculator):
621 """
622 Force-based QM/MM calculator
624 QM forces are computed using a buffer region and then mixed abruptly
625 with MM forces:
627 F^i_QMMM = { F^i_QM if i in QM region
628 { F^i_MM otherwise
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']
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
649 Parameters:
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 """
671 if len(atoms[qm_selection_mask]) == 0:
672 raise ValueError("no QM atoms selected!")
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
683 self.qm_buffer_mask = None
685 Calculator.__init__(self)
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)
698 self.qm_buffer_mask = np.zeros(len(atoms), dtype=bool)
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
705 def get_qm_cluster(self, atoms):
707 if self.qm_buffer_mask is None:
708 self.initialize_qm_buffer_mask(atoms)
710 qm_cluster = atoms[self.qm_buffer_mask]
711 del qm_cluster.constraints
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
726 if atoms.cell.orthorhombic:
727 cell_size = np.diagonal(atoms.cell)
728 else:
729 raise RuntimeError("NON-orthorhombic cell is not supported!")
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)))
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))
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)
754 qm_cluster.set_cell(Cell(np.diag(qm_cluster_cell)))
755 qm_cluster.pbc = qm_cluster_pbc
757 qm_shift = (0.5 * qm_cluster.cell.diagonal() -
758 qm_cluster.positions.mean(axis=0))
760 if 'cell_origin' in qm_cluster.info:
761 del qm_cluster.info['cell_origin']
763 # center the cluster only in non pbc directions
764 qm_cluster.positions[:, ~qm_cluster_pbc] += qm_shift[~qm_cluster_pbc]
766 return qm_cluster
768 def calculate(self, atoms, properties, system_changes):
769 Calculator.calculate(self, atoms, properties, system_changes)
771 qm_cluster = self.get_qm_cluster(atoms)
773 forces = self.mm_calc.get_forces(atoms)
774 qm_forces = self.qm_calc.get_forces(qm_cluster)
776 forces[self.qm_selection_mask] = \
777 qm_forces[self.qm_selection_mask[self.qm_buffer_mask]]
779 if self.zero_mean:
780 # Target is that: forces.sum(axis=1) == [0., 0., 0.]
781 forces[:] -= forces.mean(axis=0)
783 self.results['forces'] = forces
784 self.results['energy'] = 0.0
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
800 region = np.full_like(atoms, "MM")
802 region[self.qm_selection_mask] = (
803 np.full_like(region[self.qm_selection_mask], "QM"))
805 buffer_only_mask = self.qm_buffer_mask & ~self.qm_selection_mask
807 region[buffer_only_mask] = np.full_like(region[buffer_only_mask],
808 "buffer")
810 if print_mapping:
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}")
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}")
823 return region
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"
832 self.qm_buffer_mask = self.qm_selection_mask ^ buffer_mask
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
846 region = self.get_region_from_masks(atoms=atoms)
848 atoms_copy = atoms.copy()
849 atoms_copy.new_array("region", region)
851 atoms_copy.calc = self # to keep the calculation results
853 atoms_copy.write(filename, format='extxyz')
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
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.
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 """
878 from ase.io import read
879 atoms = read(filename, format='extxyz')
881 if "region" in atoms.arrays:
882 region = atoms.get_array("region")
883 else:
884 raise RuntimeError("Please provide extxyz file with 'region' array")
886 dummy_qm_mask = np.full_like(atoms, False, dtype=bool)
887 dummy_qm_mask[0] = True
889 self = cls(atoms, dummy_qm_mask,
890 qm_calc, mm_calc, buffer_width=1.0)
892 self.set_masks_from_region(region)
894 return self