Coverage for gpaw/electrostatic_potential.py: 28%

1from __future__ import annotations 

2 

3from math import pi, sqrt 

4from typing import TYPE_CHECKING 

5 

6import numpy as np 

7from ase.units import Bohr, Ha 

8 

9from gpaw.atom.shapefunc import shape_functions 

10from gpaw.core.arrays import DistributedArrays 

11from gpaw.core.atom_arrays import AtomArrays 

12from gpaw.core.uniform_grid import UGArray 

13from gpaw.setup import Setups 

14from gpaw.typing import Array1D, ArrayLike2D 

15from gpaw.utilities import pack_density 

16 

17if TYPE_CHECKING: 

18 from gpaw.new.calculation import DFTCalculation 

19 

20 

21class ElectrostaticPotential: 

22 def __init__(self, 

23 vHt_x: DistributedArrays, 

24 W_aL: AtomArrays, 

25 Q_aL: AtomArrays, 

26 D_asii: AtomArrays, 

27 relpos_ac: ArrayLike2D, 

28 setups: Setups): 

29 self.vHt_x = vHt_x 

30 self.W_aL = W_aL 

31 self.Q_aL = Q_aL 

32 self.D_asii = D_asii 

33 self.relpos_ac = relpos_ac 

34 self.setups = setups 

35 

36 # Caching of interpolated pseudo-potential: 

37 self._grid_spacing = -1.0 

38 self._vHt_R: UGArray | None = None 

39 

40 @classmethod 

41 def from_calculation(cls, calculation: DFTCalculation): 

42 density = calculation.density 

43 potential, _, W_aL = calculation.pot_calc.calculate(density) 

44 Q_aL = density.calculate_compensation_charge_coefficients() 

45 return cls(potential.vHt_x, 

46 W_aL, 

47 Q_aL, 

48 density.D_asii, 

49 calculation.relpos_ac, 

50 calculation.setups) 

51 

52 def atomic_potentials(self) -> Array1D: 

53 W_aL = self.W_aL.gather(broadcast=True) 

54 return W_aL.data[::9] * (Ha / (4 * pi)**0.5) 

55 

56 def atomic_corrections(self): 

57 """Calculate PAW correction to average electrostatic potential.""" 

58 dEH_a = np.zeros(len(self.setups)) 

59 for a, D_sii in self.D_asii.items(): 

60 setup = self.setups[a] 

61 D_p = pack_density(D_sii.sum(0)) 

62 dEH_a[a] = setup.dEH0 + setup.dEH_p @ D_p 

63 self.D_asii.layout.atomdist.comm.sum(dEH_a) 

64 return dEH_a * Ha * Bohr**3 

65 

66 def pseudo_potential(self, 

67 grid_spacing: float = 0.05, # Ang 

68 ) -> UGArray: 

69 return self._pseudo_potential(grid_spacing / Bohr).scaled(Bohr, Ha) 

70 

71 def _pseudo_potential(self, 

72 grid_spacing: float, # Bohr 

73 ) -> UGArray: 

74 if grid_spacing == self._grid_spacing: 

75 assert self._vHt_R is not None 

76 return self._vHt_R 

77 

78 vHt_x = self.vHt_x 

79 if isinstance(self.vHt_x, UGArray): 

80 vHt_x = self.vHt_x.to_pbc_grid() 

81 grid = vHt_x.desc.uniform_grid_with_grid_spacing(grid_spacing / Bohr) 

82 self._vHt_R = vHt_x.interpolate(grid=grid) 

83 self._grid_spacing = grid_spacing 

84 return self._vHt_R 

85 

86 def all_electron_potential(self, 

87 grid_spacing: float = 0.05, # Ang 

88 rcgauss: float = 0.02, # Ang 

89 npoints: int = 200) -> UGArray: 

90 """Interpolate electrostatic potential. 

91 

92 Return value in eV. 

93 

94 ae: bool 

95 Add PAW correction to get the all-electron potential. 

96 rcgauss: float 

97 Width of gaussian (in Angstrom) used to represent the nuclear 

98 charge. 

99 """ 

100 vHt_R = self._pseudo_potential(grid_spacing / Bohr).copy() 

101 

102 dv_a = [] 

103 for a, D_sii in self.D_asii.items(): 

104 setup = self.setups[a] 

105 c = setup.xc_correction 

106 rgd = c.rgd 

107 params = setup.data.shape_function.copy() 

108 params['lmax'] = 0 

109 ghat_g = shape_functions(rgd, **params)[0] 

110 Z_g = shape_functions(rgd, 'gauss', rcgauss, lmax=0)[0] * setup.Z 

111 D_p = pack_density(D_sii.sum(axis=0)) 

112 D_q = D_p @ c.B_pqL[:, :, 0] 

113 dn_g = D_q @ (c.n_qg - c.nt_qg) * sqrt(4 * pi) 

114 dn_g += 4 * pi * (c.nc_g - c.nct_g) 

115 dn_g -= Z_g 

116 dn_g -= self.Q_aL[a][0] * ghat_g * sqrt(4 * pi) 

117 dv_g = rgd.poisson(dn_g) / sqrt(4 * pi) 

118 dv_g[1:] /= rgd.r_g[1:] 

119 dv_g[0] = dv_g[1] 

120 dv_g[-1] = 0.0 

121 dv_a.append([rgd.spline(dv_g, points=npoints)]) 

122 

123 dv_aR = vHt_R.desc.atom_centered_functions(dv_a, self.relpos_ac) 

124 dv_aR.add_to(vHt_R) 

125 return vHt_R.scaled(Bohr, Ha)