Coverage for gpaw/solvation/cavity.py: 91%
634 statements
« prev ^ index » next coverage.py v7.7.1, created at 2025-07-09 00:21 +0000
1import numpy as np
2from ase.units import kB, Hartree, Bohr
3from ase.data.vdw import vdw_radii
5from gpaw.solvation.gridmem import NeedsGD
6from gpaw.fd_operators import Gradient
7from gpaw.io.logger import indent
9BAD_RADIUS_MESSAGE = 'All atomic radii have to be finite and >= zero.'
12def set_log_and_check_radii(obj, atoms, log):
13 radii = np.array(
14 [obj.atomic_radii.get(symbol, vdw_radii[Z])
15 for symbol, Z in zip(atoms.symbols, atoms.numbers)], dtype=float)
16 obj.atomic_radii_output = radii.copy()
17 obj.symbols = atoms.get_chemical_symbols()
18 log_radii, a_index, na = np.unique(radii, return_index=True,
19 return_counts=True)
20 log_symbols = [atoms.get_chemical_symbols()[a] for a in a_index]
21 log(f' Atomic radii for {obj.__class__.__name__}:')
22 log(' ' * 4 + 'Type' + ' ' * 4 + 'Radius' + ' ' * 3 + 'No. of atoms')
23 outstring = ''
24 for ia in range(len(log_radii)):
25 for col in [f'{log_symbols[ia]:>3}', f'{log_radii[ia]:.3f}',
26 f'{na[ia]}\n']:
27 outstring += ' ' * 5 + col
28 log(outstring)
29 if not np.isfinite(radii).all() or (radii < 0).any():
30 raise ValueError(BAD_RADIUS_MESSAGE)
33def get_pbc_positions(atoms, r_max):
34 """Return dict mapping atom index to positions in Bohr.
36 With periodic boundary conditions, it also includes neighbouring
37 cells up to a distance of r_max (in Bohr).
38 """
39 # code snippet taken from ase/calculators/vdwcorrection.py
40 pbc_c = atoms.get_pbc()
41 cell_cv = atoms.get_cell() / Bohr
42 Rcell_c = np.sqrt(np.sum(cell_cv ** 2, axis=1))
43 ncells_c = np.ceil(np.where(pbc_c, 1. + r_max / Rcell_c, 1)).astype(int)
44 pos_aav = {}
45 # loop over all atoms in the cell (and neighbour cells for PBC)
46 for index1, atom in enumerate(atoms):
47 pos = atom.position / Bohr
48 pos_aav[index1] = np.empty((np.prod(ncells_c * 2 - 1), 3))
49 # loops over neighbour cells
50 index2 = 0
51 for ix in range(-ncells_c[0] + 1, ncells_c[0]):
52 for iy in range(-ncells_c[1] + 1, ncells_c[1]):
53 for iz in range(-ncells_c[2] + 1, ncells_c[2]):
54 i_c = np.array([ix, iy, iz])
55 pos_aav[index1][index2, :] = pos + np.dot(i_c, cell_cv)
56 index2 += 1
57 return pos_aav
60class Cavity(NeedsGD):
61 """Base class for representing a cavity in the solvent.
63 Attributes:
64 g_g -- The cavity on the fine grid. It varies from zero at the
65 solute location to one in the bulk solvent.
66 del_g_del_n_g -- The partial derivative of the cavity with respect to
67 the electron density on the fine grid.
68 grad_g_vg -- The gradient of the cavity on the fine grid.
69 V -- global Volume in Bohr ** 3 or None
70 A -- global Area in Bohr ** 2 or None
71 """
73 def __init__(self, surface_calculator=None, volume_calculator=None):
74 """Constructor for the Cavity class.
76 Arguments:
77 surface_calculator -- A SurfaceCalculator instance or None
78 volume_calculator -- A VolumeCalculator instance or None
79 """
80 NeedsGD.__init__(self)
81 self.g_g = None
82 self.del_g_del_n_g = None
83 self.grad_g_vg = None
84 if isinstance(surface_calculator, dict):
85 surface_calculator = GradientSurface(**surface_calculator)
86 if isinstance(volume_calculator, dict):
87 volume_calculator = KB51Volume(**volume_calculator)
88 self.surface_calculator = surface_calculator
89 self.volume_calculator = volume_calculator
90 self.V = None # global Volume
91 self.A = None # global Surface
93 def todict(self):
94 dct = {}
95 if self.surface_calculator is not None:
96 dct['surface_calculator'] = self.surface_calculator.todict()
97 if self.volume_calculator is not None:
98 dct['volume_calculator'] = self.volume_calculator.todict()
99 return dct
101 @classmethod
102 def from_dict(cls, dct):
103 if not isinstance(dct, dict):
104 return dct
105 return EffectivePotentialCavity(**dct)
107 def write(self, writer):
108 pass
110 def read(self, reader):
111 pass
113 def estimate_memory(self, mem):
114 ngrids = 1 + self.depends_on_el_density
115 mem.subnode('Distribution Function', ngrids * self.gd.bytecount())
116 mem.subnode(
117 'Gradient of Distribution Function', 3 * self.gd.bytecount()
118 )
119 if self.surface_calculator is not None:
120 self.surface_calculator.estimate_memory(
121 mem.subnode('Surface Calculator')
122 )
123 if self.volume_calculator is not None:
124 self.volume_calculator.estimate_memory(
125 mem.subnode('Volume Calculator')
126 )
128 def update(self, atoms, density):
129 """Update the cavity and its gradient.
131 atoms are None, iff they have not changed.
133 Return whether the cavity has changed.
134 """
135 raise NotImplementedError()
137 def update_vol_surf(self):
138 """Update volume and surface."""
139 if self.surface_calculator is not None:
140 self.surface_calculator.update(self)
141 if self.volume_calculator is not None:
142 self.volume_calculator.update(self)
144 def communicate_vol_surf(self, world):
145 """Communicate global volume and surface."""
146 if self.surface_calculator is not None:
147 A = np.array([self.surface_calculator.A])
148 self.gd.comm.sum(A)
149 world.broadcast(A, 0)
150 self.A = A[0]
151 else:
152 self.A = None
153 if self.volume_calculator is not None:
154 V = np.array([self.volume_calculator.V])
155 self.gd.comm.sum(V)
156 world.broadcast(V, 0)
157 self.V = V[0]
158 else:
159 self.V = None
161 def allocate(self):
162 NeedsGD.allocate(self)
163 self.g_g = self.gd.empty()
164 self.grad_g_vg = self.gd.empty(3)
165 if self.depends_on_el_density:
166 self.del_g_del_n_g = self.gd.empty()
167 if self.surface_calculator is not None:
168 self.surface_calculator.allocate()
169 if self.volume_calculator is not None:
170 self.volume_calculator.allocate()
172 def set_grid_descriptor(self, gd):
173 NeedsGD.set_grid_descriptor(self, gd)
174 if self.surface_calculator is not None:
175 self.surface_calculator.set_grid_descriptor(gd)
176 if self.volume_calculator is not None:
177 self.volume_calculator.set_grid_descriptor(gd)
179 def get_del_r_vg(self, atom_index, density):
180 """Return spatial derivatives with respect to atomic position."""
181 raise NotImplementedError()
183 @property
184 def depends_on_el_density(self):
185 """Return whether the cavity depends on the electron density."""
186 raise NotImplementedError()
188 @property
189 def depends_on_atomic_positions(self):
190 """Return whether the cavity depends explicitly on atomic positions."""
191 raise NotImplementedError()
193 def __str__(self):
194 s = f'Cavity: {self.__class__.__name__}\n'
195 for calc, calcname in ((self.surface_calculator, 'Surface'),
196 (self.volume_calculator, 'Volume')):
197 s += f' {calcname} Calculator: '
198 if calc is None:
199 s += 'None\n'
200 else:
201 s += str(calc)
202 return s
204 def update_atoms(self, atoms, log):
205 """Inexpensive initialization when atoms change."""
206 pass
208 def summary(self, log):
209 """Log cavity surface area and volume."""
210 A = (f'{self.A * Bohr ** 2:.5f}' if self.A is not None
211 else 'not calculated (no calculator defined)')
212 V = (f'{self.V * Bohr ** 3:.5f}' if self.V is not None
213 else 'not calculated (no calculator defined)')
214 log(f'Solvation cavity surface area: {A}')
215 log(f'Solvation cavity volume: {V}')
218class EffectivePotentialCavity(Cavity):
219 """Cavity built from effective potential and Boltzmann distribution.
221 See also
222 A. Held and M. Walter, J. Chem. Phys. 141, 174108 (2014).
223 """
225 def __init__(self,
226 effective_potential,
227 temperature,
228 surface_calculator=None,
229 volume_calculator=None):
230 """Constructor for the EffectivePotentialCavity class.
232 Additional arguments not present in base Cavity class:
233 effective_potential -- A Potential instance.
234 temperature -- Temperature for the Boltzmann distribution
235 in Kelvin.
236 """
237 Cavity.__init__(self, surface_calculator, volume_calculator)
238 self.effective_potential = Potential.from_dict(effective_potential)
239 self.temperature = float(temperature)
240 self.minus_beta = -1. / (kB * temperature / Hartree)
242 def todict(self):
243 return {
244 'effective_potential': {
245 'name': self.effective_potential.__class__.__name__,
246 **self.effective_potential.todict()},
247 'temperature': self.temperature,
248 **super().todict()}
250 def write(self, writer):
251 writer.write(effective_potential=self.effective_potential,
252 temperature=self.temperature,
253 surface_calculator=self.surface_calculator,
254 volume_calculator=self.volume_calculator)
256 # For future, not used at the moment
257 def read(self, reader):
258 c = reader.parameters.cavity
259 self.effective_potential = c.effective_potential
260 self.temperature = c.temperature
261 self.surface_calculator = c.surface_calculator
262 self.volume_calculator = c.volume_calculator
264 def estimate_memory(self, mem):
265 Cavity.estimate_memory(self, mem)
266 self.effective_potential.estimate_memory(
267 mem.subnode('Effective Potential')
268 )
270 def set_grid_descriptor(self, gd):
271 Cavity.set_grid_descriptor(self, gd)
272 self.effective_potential.set_grid_descriptor(gd)
274 def allocate(self):
275 Cavity.allocate(self)
276 self.effective_potential.allocate()
278 def update(self, atoms, density):
279 if not self.effective_potential.update(atoms, density):
280 return False
281 u_g = self.effective_potential.u_g
282 np.exp(u_g * self.minus_beta, out=self.g_g)
283 if self.depends_on_el_density:
284 self.del_g_del_n_g.fill(self.minus_beta)
285 self.del_g_del_n_g *= self.g_g
286 self.del_g_del_n_g *= self.effective_potential.del_u_del_n_g
287 # This class supports the GradientSurface API,
288 # so we can use it for the gradient also
289 grad_inner_vg = self.get_grad_inner()
290 del_outer_del_inner = self.get_del_outer_del_inner()
291 for i in (0, 1, 2):
292 np.multiply(
293 grad_inner_vg[i],
294 del_outer_del_inner,
295 self.grad_g_vg[i]
296 )
297 return True
299 def get_del_r_vg(self, atom_index, density):
300 u = self.effective_potential
301 del_u_del_r_vg = u.get_del_r_vg(atom_index, density)
302 # there should be no more NaNs now, but let's keep the hint
303 # asserts lim_(||r - r_atom|| -> 0) dg/du * du/dr_atom = 0
304 # del_u_del_r_vg[np.isnan(del_u_del_r_vg)] = .0
305 return self.minus_beta * self.g_g * del_u_del_r_vg
307 @property
308 def depends_on_el_density(self):
309 return self.effective_potential.depends_on_el_density
311 @property
312 def depends_on_atomic_positions(self):
313 return self.effective_potential.depends_on_atomic_positions
315 def __str__(self):
316 s = Cavity.__str__(self)
317 s += f' {self.__class__.__name__}\n'
318 s += indent(f'{self.effective_potential}')
319 s += indent(f' temperature: {self.temperature}K')
320 return s
322 def update_atoms(self, atoms, log):
323 self.effective_potential.update_atoms(atoms, log)
325 # --- BEGIN GradientSurface API ---
327 def get_grad_inner(self):
328 return self.effective_potential.grad_u_vg
330 def get_del_outer_del_inner(self):
331 return self.minus_beta * self.g_g
333 # --- END GradientSurface API ---
336class Potential(NeedsGD):
337 """Base class for describing an effective potential.
339 Attributes:
340 u_g -- The potential on the fine grid in Hartree.
341 grad_u_vg -- The gradient of the potential of the fine grid.
342 del_u_del_n_g -- Partial derivative with respect to the electron density.
343 """
345 def __init__(self):
346 NeedsGD.__init__(self)
347 self.u_g = None
348 self.del_u_del_n_g = None
349 self.grad_u_vg = None
351 @classmethod
352 def from_dict(self, dct):
353 if not isinstance(dct, dict):
354 return dct
355 dct = dct.copy()
356 name = dct.pop('name')
357 if name == 'Power12Potential':
358 return Power12Potential(**dct)
359 assert name == 'SJMPower12Potential'
360 from gpaw.solvation.sjm import SJMPower12Potential
361 return SJMPower12Potential(**dct)
363 @property
364 def depends_on_el_density(self):
365 """Return whether the cavity depends on the electron density."""
366 raise NotImplementedError()
368 @property
369 def depends_on_atomic_positions(self):
370 """returns whether the cavity depends explicitly on atomic positions"""
371 raise NotImplementedError()
373 def estimate_memory(self, mem):
374 ngrids = 1 + self.depends_on_el_density
375 mem.subnode('Potential', ngrids * self.gd.bytecount())
376 mem.subnode('Gradient of Potential', 3 * self.gd.bytecount())
378 def allocate(self):
379 NeedsGD.allocate(self)
380 self.u_g = self.gd.empty()
381 if self.depends_on_el_density:
382 self.del_u_del_n_g = self.gd.empty()
383 self.grad_u_vg = self.gd.empty(3)
385 def update(self, atoms, density):
386 """Update the potential and its gradient.
388 atoms are None, iff they have not changed.
390 Return whether the potential has changed.
391 """
392 raise NotImplementedError()
394 def get_del_r_vg(self, atom_index, density):
395 """Return spatial derivatives with respect to atomic position."""
396 raise NotImplementedError()
398 def __str__(self):
399 return f' Potential: {self.__class__.__name__}\n'
401 def update_atoms(self, atoms, log):
402 """Inexpensive initialization when atoms change."""
403 pass
406def get_vdw_radii(atoms):
407 """Returns a list of van der Waals radii for a given atoms object."""
408 return [vdw_radii[n] for n in atoms.numbers]
411class Power12Potential(Potential):
412 """Inverse power law potential.
414 A 1 / r ** 12 repulsive potential taking the value u0 at the atomic radius.
416 See also
417 A. Held and M. Walter, J. Chem. Phys. 141, 174108 (2014).
419 Parameters:
421 atomic_radii: function
422 Callable mapping an ase.Atoms object to an iterable of atomic radii
423 in Angstroms. If not provided, defaults to van der Waals radii.
424 u0: float
425 Strength of the potential at the atomic radius in eV.
426 Defaults to 0.18 eV, the best-fit value for water from Held &
427 Walter.
428 pbc_cutoff: float
429 Cutoff in eV for including neighbor cells in a calculation with
430 periodic boundary conditions.
431 """
432 depends_on_el_density = False
433 depends_on_atomic_positions = True
435 def __init__(self, atomic_radii=None, u0=0.180, pbc_cutoff=1e-6,
436 tiny=1e-10):
437 Potential.__init__(self)
438 self.atomic_radii = atomic_radii or {}
439 self.u0 = float(u0)
440 self.pbc_cutoff = float(pbc_cutoff)
441 self.tiny = float(tiny)
442 self.r12_a = None
443 self.r_vg = None
444 self.pos_aav = None
445 self.del_u_del_r_vg = None
446 self.atomic_radii_output = None
447 self.symbols = None
449 def todict(self):
450 return {
451 'atomic_radii': self.atomic_radii,
452 'u0': self.u0,
453 'pbc_cutoff': self.pbc_cutoff,
454 'tiny': self.tiny}
456 def estimate_memory(self, mem):
457 Potential.estimate_memory(self, mem)
458 nbytes = self.gd.bytecount()
459 mem.subnode('Coordinates', 3 * nbytes)
460 mem.subnode('Atomic Position Derivative', 3 * nbytes)
462 def allocate(self):
463 Potential.allocate(self)
464 self.r_vg = self.gd.get_grid_point_coordinates()
465 self.del_u_del_r_vg = self.gd.empty(3)
467 def update(self, atoms, density):
468 if atoms is None:
469 return False
470 self.r12_a = (self.atomic_radii_output / Bohr) ** 12
471 r_cutoff = (self.r12_a.max() * self.u0 / self.pbc_cutoff) ** (1. / 12.)
472 self.pos_aav = get_pbc_positions(atoms, r_cutoff)
473 self.u_g.fill(.0)
474 self.grad_u_vg.fill(.0)
475 na = np.newaxis
476 for index, pos_av in self.pos_aav.items():
477 r12 = self.r12_a[index]
478 for pos_v in pos_av:
479 origin_vg = pos_v[:, na, na, na]
480 r_diff_vg = self.r_vg - origin_vg
481 r_diff2_g = (r_diff_vg ** 2).sum(0)
482 r_diff2_g[r_diff2_g < self.tiny] = self.tiny
483 u_g = r12 / r_diff2_g ** 6
484 self.u_g += u_g
485 u_g /= r_diff2_g
486 r_diff_vg *= u_g[na, ...]
487 self.grad_u_vg += r_diff_vg
488 self.u_g *= self.u0 / Hartree
489 self.grad_u_vg *= -12. * self.u0 / Hartree
490 # avoid overflow in norm calculation:
491 self.grad_u_vg[self.grad_u_vg < -1e20] = -1e20
492 self.grad_u_vg[self.grad_u_vg > 1e20] = 1e20
493 return True
495 def get_del_r_vg(self, atom_index, density):
496 u0 = self.u0 / Hartree
497 r12 = self.r12_a[atom_index]
498 na = np.newaxis
499 self.del_u_del_r_vg.fill(.0)
500 for pos_v in self.pos_aav[atom_index]:
501 origin_vg = pos_v[:, na, na, na]
502 diff_vg = self.r_vg - origin_vg
503 diff2_g = (diff_vg ** 2).sum(0)
504 diff2_g[diff2_g < self.tiny] = self.tiny
505 diff2_g **= 7
506 diff_vg /= diff2_g[na, ...]
507 self.del_u_del_r_vg += diff_vg
508 self.del_u_del_r_vg *= (12. * u0 * r12)
509 return self.del_u_del_r_vg
511 def __str__(self):
512 s = Potential.__str__(self)
513 s += indent(f' u0: {self.u0}eV\n')
514 s += indent(f' pbc_cutoff: {self.pbc_cutoff}\n')
515 s += indent(f' tiny: {self.tiny}\n')
516 return s
518 def update_atoms(self, atoms, log):
519 set_log_and_check_radii(self, atoms, log)
521 def write(self, writer):
522 writer.write(
523 name=self.__class__.__name__,
524 atomic_radii=self.atomic_radii_output,
525 u0=self.u0,
526 pbc_cutoff=self.pbc_cutoff,
527 tiny=self.tiny)
530class SmoothStepCavity(Cavity):
531 """Base class for cavities based on a smooth step function and a density.
533 Attributes:
534 del_g_del_rho_g -- Partial derivative with respect to the density.
535 """
537 def __init__(
538 self,
539 density,
540 surface_calculator=None, volume_calculator=None,
541 ):
542 """Constructor for the SmoothStepCavity class.
544 Additional arguments not present in the base Cavity class:
545 density -- A Density instance
546 """
547 Cavity.__init__(self, surface_calculator, volume_calculator)
548 self.del_g_del_rho_g = None
549 self.density = density
551 @property
552 def depends_on_el_density(self):
553 return self.density.depends_on_el_density
555 @property
556 def depends_on_atomic_positions(self):
557 return self.density.depends_on_atomic_positions
559 def set_grid_descriptor(self, gd):
560 Cavity.set_grid_descriptor(self, gd)
561 self.density.set_grid_descriptor(gd)
563 def estimate_memory(self, mem):
564 Cavity.estimate_memory(self, mem)
565 mem.subnode('Cavity Derivative', self.gd.bytecount())
566 self.density.estimate_memory(mem.subnode('Density'))
568 def allocate(self):
569 Cavity.allocate(self)
570 self.del_g_del_rho_g = self.gd.empty()
571 self.density.allocate()
573 def update(self, atoms, density):
574 if not self.density.update(atoms, density):
575 return False
576 self.update_smooth_step(self.density.rho_g)
577 if self.depends_on_el_density:
578 np.multiply(
579 self.del_g_del_rho_g,
580 self.density.del_rho_del_n_g,
581 self.del_g_del_n_g
582 )
583 # This class supports the GradientSurface API,
584 # so we can use it for the gradient also
585 grad_inner_vg = self.get_grad_inner()
586 del_outer_del_inner = self.get_del_outer_del_inner()
587 for i in (0, 1, 2):
588 np.multiply(
589 grad_inner_vg[i],
590 del_outer_del_inner,
591 self.grad_g_vg[i])
592 return True
594 def update_smooth_step(self, rho_g):
595 """Calculate self.g_g and self.del_g_del_rho_g."""
596 raise NotImplementedError()
598 def get_del_r_vg(self, atom_index, density):
599 return self.del_g_del_rho_g * self.density.get_del_r_vg(
600 atom_index, density
601 )
603 def __str__(self):
604 s = Cavity.__str__(self)
605 s += indent(str(self.density))
606 return s
608 def update_atoms(self, atoms, log):
609 self.density.update_atoms(atoms, log)
611 # --- BEGIN GradientSurface API ---
612 def get_grad_inner(self):
613 return self.density.grad_rho_vg
615 def get_del_outer_del_inner(self):
616 return self.del_g_del_rho_g
618 # --- END GradientSurface API ---
621class Density(NeedsGD):
622 """Class representing a density for the use with a SmoothStepCavity.
624 Attributes:
625 rho_g -- The density on the fine grid in 1 / Bohr ** 3.
626 del_rho_del_n_g -- Partial derivative with respect to the electron density.
627 grad_rho_vg -- The gradient of the density on the fine grid.
628 """
630 def __init__(self):
631 NeedsGD.__init__(self)
632 self.rho_g = None
633 self.del_rho_del_n_g = None
634 self.grad_rho_vg = None
636 def estimate_memory(self, mem):
637 nbytes = self.gd.bytecount()
638 mem.subnode('Density', nbytes)
639 if self.depends_on_el_density:
640 mem.subnode('Density Derivative', nbytes)
641 mem.subnode('Gradient of Density', 3 * nbytes)
643 def allocate(self):
644 NeedsGD.allocate(self)
645 self.rho_g = self.gd.empty()
646 if self.depends_on_el_density:
647 self.del_rho_del_n_g = self.gd.empty()
648 self.grad_rho_vg = self.gd.empty(3)
650 def update(self, atoms, density):
651 raise NotImplementedError()
653 @property
654 def depends_on_el_density(self):
655 raise NotImplementedError()
657 @property
658 def depends_on_atomic_positions(self):
659 raise NotImplementedError()
661 def __str__(self):
662 return f"Density: {self.__class__.__name__}\n"
664 def update_atoms(self, atoms, log):
665 """Inexpensive initialization when atoms change."""
666 pass
669class FDGradientDensity(Density):
670 """Base class for all Density classes with finite difference gradient"""
672 def __init__(self, boundary_value, nn=3):
673 """Constructur for the FDGradientDensity class.
675 Arguments:
676 boundary_value -- Boundary value of rho_g
677 (for non-periodic directions).
678 nn -- Stencil size for the finite difference gradient.
679 """
680 Density.__init__(self)
681 self.boundary_value = boundary_value
682 self.nn = nn
683 self.gradient = None
685 def allocate(self):
686 Density.allocate(self)
687 self.gradient = [
688 Gradient(self.gd, i, 1.0, self.nn) for i in (0, 1, 2)
689 ]
691 def update(self, atoms, density):
692 changed = self.update_only_density(atoms, density)
693 if changed:
694 self.update_gradient()
695 return changed
697 def update_gradient(self):
698 if self.boundary_value != 0:
699 in_g = self.rho_g - self.boundary_value
700 else:
701 in_g = self.rho_g
702 for i in (0, 1, 2):
703 self.gradient[i].apply(in_g, self.grad_rho_vg[i])
706class ElDensity(FDGradientDensity):
707 """Wrapper class for using the electron density in a SmoothStepCavity.
709 (Hopefully small) negative values of the electron density are set to zero.
710 """
712 depends_on_el_density = True
713 depends_on_atomic_positions = False
715 def __init__(self, nn=3):
716 """Constructor for the ElDensity class.
718 Arguments:
719 nn -- Stencil size for the finite difference gradient.
720 """
721 FDGradientDensity.__init__(self, boundary_value=.0, nn=nn)
723 def allocate(self):
724 FDGradientDensity.allocate(self)
725 self.del_rho_del_n_g = 1. # free array
727 def update_only_density(self, atoms, density):
728 self.rho_g[:] = density.nt_g
729 self.rho_g[self.rho_g < .0] = .0
730 return True
733# TODO: implement analytic gradient for SSS09Density
734class SSS09Density(FDGradientDensity):
735 """Fake density from atomic radii for the use in a SmoothStepCavity.
737 Following V. M. Sanchez, M. Sued and D. A. Scherlis,
738 J. Chem. Phys. 131, 174108 (2009).
739 """
741 depends_on_el_density = False
742 depends_on_atomic_positions = True
744 def __init__(self, atomic_radii, pbc_cutoff=1e-3, nn=3, tiny=1e-100):
745 """Constructor for the SSS09Density class.
747 Arguments:
748 atomic_radii -- Callable mapping an ase.Atoms object
749 to an iterable of atomic radii in Angstroms.
750 pbc_cutoff -- Cutoff in eV for including neighbor cells in
751 a calculation with periodic boundary conditions.
752 nn -- Stencil size for the finite difference gradient.
753 """
754 FDGradientDensity.__init__(self, boundary_value=.0, nn=nn)
755 self.atomic_radii = atomic_radii
756 self.atomic_radii_output = None
757 self.symbols = None
758 self.pbc_cutoff = float(pbc_cutoff)
759 self.tiny = float(tiny)
760 self.pos_aav = None
761 self.r_vg = None
762 self.del_rho_del_r_vg = None
764 def estimate_memory(self, mem):
765 FDGradientDensity.estimate_memory(self, mem)
766 nbytes = self.gd.bytecount()
767 mem.subnode('Coordinates', 3 * nbytes)
768 mem.subnode('Atomic Position Derivative', 3 * nbytes)
770 def allocate(self):
771 FDGradientDensity.allocate(self)
772 self.r_vg = self.gd.get_grid_point_coordinates()
773 self.del_rho_del_r_vg = self.gd.empty(3)
775 def update_only_density(self, atoms, density):
776 if atoms is None:
777 return False
778 r_a = self.atomic_radii_output / Bohr
779 r_cutoff = r_a.max() - np.log(self.pbc_cutoff)
780 self.pos_aav = get_pbc_positions(atoms, r_cutoff)
781 self.rho_g.fill(.0)
782 na = np.newaxis
783 for index, pos_av in self.pos_aav.items():
784 for pos_v in pos_av:
785 origin_vg = pos_v[:, na, na, na]
786 r_diff_vg = self.r_vg - origin_vg
787 norm_r_diff_g = (r_diff_vg ** 2).sum(0) ** .5
788 self.rho_g += np.exp(r_a[index] - norm_r_diff_g)
789 return True
791 def get_del_r_vg(self, atom_index, density):
792 r_a = self.atomic_radii_output[atom_index] / Bohr
793 na = np.newaxis
794 self.del_rho_del_r_vg.fill(.0)
795 for pos_v in self.pos_aav[atom_index]:
796 origin_vg = pos_v[:, na, na, na]
797 r_diff_vg = self.r_vg - origin_vg
798 norm_r_diff_g = np.sum(r_diff_vg ** 2, axis=0) ** .5
799 norm_r_diff_g[norm_r_diff_g < self.tiny] = self.tiny
800 exponential = np.exp(r_a - norm_r_diff_g)
801 exponential /= norm_r_diff_g
802 r_diff_vg *= exponential[na, ...]
803 self.del_rho_del_r_vg += r_diff_vg
804 return self.del_rho_del_r_vg
806 def __str__(self):
807 s = FDGradientDensity.__str__(self)
808 s += indent(f' pbc_cutoff: {self.pbc_cutoff}\n')
809 s += indent(f' tiny: {self.tiny}\n')
810 return s
812 def update_atoms(self, atoms, log):
813 set_log_and_check_radii(self, atoms, log)
816class ADM12SmoothStepCavity(SmoothStepCavity):
817 """Cavity from smooth step function.
819 Following O. Andreussi, I. Dabo, and N. Marzari,
820 J. Chem. Phys. 136, 064102 (2012).
821 """
823 def __init__(
824 self,
825 rhomin, rhomax, epsinf,
826 density,
827 surface_calculator=None, volume_calculator=None
828 ):
829 """Constructor for the ADM12SmoothStepCavity class.
831 Additional arguments not present in the SmoothStepCavity class:
832 rhomin -- Lower density isovalue in 1 / Angstrom ** 3.
833 rhomax -- Upper density isovalue in 1 / Angstrom ** 3.
834 epsinf -- Static dielectric constant of the solvent.
835 """
836 SmoothStepCavity.__init__(
837 self, density, surface_calculator, volume_calculator
838 )
839 self.rhomin = float(rhomin)
840 self.rhomax = float(rhomax)
841 self.epsinf = float(epsinf)
843 def update_smooth_step(self, rho_g):
844 eps = self.epsinf
845 inside = rho_g > self.rhomax * Bohr ** 3
846 outside = rho_g < self.rhomin * Bohr ** 3
847 transition = np.logical_not(
848 np.logical_or(inside, outside)
849 )
850 self.g_g[inside] = .0
851 self.g_g[outside] = 1.
852 self.del_g_del_rho_g.fill(.0)
853 t, dt = self._get_t_dt(np.log(rho_g[transition]))
854 if eps == 1.0:
855 # lim_{eps -> 1} (eps - eps ** t) / (eps - 1) = 1 - t
856 self.g_g[transition] = t
857 self.del_g_del_rho_g[transition] = dt / rho_g[transition]
858 else:
859 eps_to_t = eps ** t
860 self.g_g[transition] = (eps_to_t - 1.) / (eps - 1.)
861 self.del_g_del_rho_g[transition] = (
862 eps_to_t * np.log(eps) * dt
863 ) / (
864 rho_g[transition] * (eps - 1.)
865 )
867 def _get_t_dt(self, x):
868 lnmax = np.log(self.rhomax * Bohr ** 3)
869 lnmin = np.log(self.rhomin * Bohr ** 3)
870 twopi = 2. * np.pi
871 arg = twopi * (lnmax - x) / (lnmax - lnmin)
872 t = (arg - np.sin(arg)) / twopi
873 dt = -2. * np.sin(arg / 2.) ** 2 / (lnmax - lnmin)
874 return (t, dt)
876 def __str__(self):
877 s = SmoothStepCavity.__str__(self)
878 s += indent(f' rhomin: {self.rhomin}\n')
879 s += indent(f' rhomax: {self.rhomax}\n')
880 s += indent(f' epsinf: {self.epsinf}')
881 return s
884class FG02SmoothStepCavity(SmoothStepCavity):
885 """Cavity from smooth step function.
887 Following J. Fattebert and F. Gygi, J. Comput. Chem. 23, 662 (2002).
888 """
890 def __init__(
891 self,
892 rho0,
893 beta,
894 density,
895 surface_calculator=None, volume_calculator=None
896 ):
897 """Constructor for the FG02SmoothStepCavity class.
899 Additional arguments not present in the SmoothStepCavity class:
900 rho0 -- Density isovalue in 1 / Angstrom ** 3.
901 beta -- Parameter controlling the steepness of the transition.
902 """
903 SmoothStepCavity.__init__(
904 self, density, surface_calculator, volume_calculator
905 )
906 self.rho0 = float(rho0)
907 self.beta = float(beta)
909 def update_smooth_step(self, rho_g):
910 rho0 = self.rho0 / (1. / Bohr ** 3)
911 rho_scaled_g = rho_g / rho0
912 exponent = 2. * self.beta
913 np.divide(1., rho_scaled_g ** exponent + 1., self.g_g)
914 np.multiply(
915 (-exponent / rho0) * rho_scaled_g ** (exponent - 1.),
916 self.g_g ** 2,
917 self.del_g_del_rho_g
918 )
920 def __str__(self):
921 s = SmoothStepCavity.__str__(self)
922 s += indent(f' rho0: {self.rho0}\n')
923 s += indent(f' beta: {self.beta}')
924 return s
927class SurfaceCalculator(NeedsGD):
928 """Base class for surface calculators.
930 Attributes:
931 A -- Local area in Bohr ** 2.
932 delta_A_delta_g_g -- Functional derivative with respect to the cavity
933 on the fine grid.
934 """
936 def __init__(self):
937 NeedsGD.__init__(self)
938 self.A = None
939 self.delta_A_delta_g_g = None
941 def write(self, writer):
942 pass
944 def read(self, reader):
945 pass
947 def estimate_memory(self, mem):
948 mem.subnode('Functional Derivative', self.gd.bytecount())
950 def allocate(self):
951 NeedsGD.allocate(self)
952 self.delta_A_delta_g_g = self.gd.empty()
954 def __str__(self):
955 return f'Surface Calculator: {self.__class__.__name__}\n'
957 def update(self, cavity):
958 """Calculate A and delta_A_delta_g_g."""
959 raise NotImplementedError()
962class GradientSurface(SurfaceCalculator):
963 """Class for getting the surface area from the gradient of the cavity.
965 See also W. Im, D. Beglov and B. Roux,
966 Comput. Phys. Commun. 111, 59 (1998)
967 and
968 A. Held and M. Walter, J. Chem. Phys. 141, 174108 (2014).
969 """
971 def __init__(self, nn=3):
972 SurfaceCalculator.__init__(self)
973 self.nn = nn
974 self.gradient = None
975 self.gradient_out = None
976 self.norm_grad_out = None
977 self.div_tmp = None
979 def todict(self):
980 return {'nn': self.nn}
982 def write(self, writer):
983 writer.write(
984 name='GradientSurface',
985 nn=self.nn)
987 def __str__(self):
988 return f'GradientSurface with {self.nn} nn\n'
990 def read(self, reader):
991 self.nn = reader.parameters.cavity.nn
993 def estimate_memory(self, mem):
994 SurfaceCalculator.estimate_memory(self, mem)
995 nbytes = self.gd.bytecount()
996 mem.subnode('Gradient', 4 * nbytes)
997 mem.subnode('Divergence', nbytes)
999 def allocate(self):
1000 SurfaceCalculator.allocate(self)
1001 self.gradient = [
1002 Gradient(self.gd, i, 1.0, self.nn) for i in (0, 1, 2)
1003 ]
1004 self.gradient_out = self.gd.empty(3)
1005 self.norm_grad_out = self.gd.empty()
1006 self.div_tmp = self.gd.empty()
1008 def update(self, cavity):
1009 del_outer_del_inner = cavity.get_del_outer_del_inner()
1010 sign = np.sign(del_outer_del_inner.max() + del_outer_del_inner.min())
1011 self.gradient_out[...] = cavity.get_grad_inner()
1012 self.norm_grad_out = (self.gradient_out ** 2).sum(0) ** .5
1013 self.A = sign * self.gd.integrate(
1014 del_outer_del_inner * self.norm_grad_out, global_integral=False
1015 )
1016 mask = self.norm_grad_out > 1e-12 # avoid division by zero or overflow
1017 imask = np.logical_not(mask)
1018 masked_norm_grad = self.norm_grad_out[mask]
1019 for i in (0, 1, 2):
1020 self.gradient_out[i][mask] /= masked_norm_grad
1021 # set limit for later calculations:
1022 self.gradient_out[i][imask] = .0
1023 self.calc_div(self.gradient_out, self.delta_A_delta_g_g)
1024 if sign == 1:
1025 self.delta_A_delta_g_g *= -1.
1027 def calc_div(self, vec, out):
1028 self.gradient[0].apply(vec[0], out)
1029 self.gradient[1].apply(vec[1], self.div_tmp)
1030 out += self.div_tmp
1031 self.gradient[2].apply(vec[2], self.div_tmp)
1032 out += self.div_tmp
1035class VolumeCalculator(NeedsGD):
1036 """Base class for volume calculators.
1038 Attributes:
1039 V -- Local volume in Bohr ** 3.
1040 delta_V_delta_g_g -- Functional derivative with respect to the cavity
1041 on the fine grid.
1042 """
1044 def __init__(self):
1045 NeedsGD.__init__(self)
1046 self.V = None
1047 self.delta_V_delta_g_g = None
1049 def estimate_memory(self, mem):
1050 mem.subnode('Functional Derivative', self.gd.bytecount())
1052 def allocate(self):
1053 NeedsGD.allocate(self)
1054 self.delta_V_delta_g_g = self.gd.empty()
1056 def __str__(self):
1057 return f"{self.__class__.__name__}\n"
1059 def update(self, cavity):
1060 """Calculate V and delta_V_delta_g_g"""
1061 raise NotImplementedError()
1064class KB51Volume(VolumeCalculator):
1065 """KB51 Volume Calculator.
1067 V = Integral(1 - g) + kappa_T * k_B * T
1069 Following J. G. Kirkwood and F. P. Buff, J. Chem. Phys. 19, 6, 774 (1951).
1070 """
1072 def __init__(self, compressibility=.0, temperature=.0):
1073 """Constructor for KB51Volume class.
1075 Arguments:
1076 compressibility -- Isothermal compressibility of the solvent
1077 in Angstrom ** 3 / eV.
1078 temperature -- Temperature in Kelvin.
1079 """
1080 VolumeCalculator.__init__(self)
1081 self.compressibility = float(compressibility)
1082 self.temperature = float(temperature)
1084 def todict(self):
1085 return {'compressibility': self.compressibility,
1086 'temperature': self.temperature}
1088 def __str__(self):
1089 s = VolumeCalculator.__str__(self)
1090 s += indent(f' compressibility: {self.compressibility}\n')
1091 s += indent(f' temperature: {self.temperature}\n')
1092 return s
1094 def allocate(self):
1095 VolumeCalculator.allocate(self)
1096 self.delta_V_delta_g_g = -1. # frees array
1098 def update(self, cavity):
1099 self.V = self.gd.integrate(1. - cavity.g_g, global_integral=False)
1100 V_compress = self.compressibility * kB * self.temperature / Bohr ** 3
1101 self.V += V_compress / self.gd.comm.size