Coverage for gpaw/pw/poisson.py: 100%

1from math import pi 

2 

3import numpy as np 

4from ase.utils import seterr 

5from scipy.special import erf 

6 

7from gpaw.pw.density import ReciprocalSpaceDensity 

8from gpaw.pw.descriptor import PWDescriptor 

9from gpaw.typing import Array1D 

10 

11 

12class ReciprocalSpacePoissonSolver: 

13 def __init__(self, 

14 pd: PWDescriptor, 

15 charge: float = 0.0, 

16 strength: float = 1.0): 

17 self.pd = pd 

18 self.charge = charge 

19 self.strength = strength 

20 

21 self.G2_q = pd.G2_qG[0].copy() 

22 if pd.gd.comm.rank == 0: 

23 # Avoid division by zero: 

24 self.G2_q[0] = 1.0 

25 

26 def __str__(self): 

27 return f'Uniform background charge: {self.charge:.3f} electrons' 

28 

29 def estimate_memory(self, mem): 

30 pass 

31 

32 def solve(self, 

33 vHt_q: Array1D, 

34 dens: ReciprocalSpaceDensity) -> float: 

35 """Solve Poisson equeation. 

36 

37 Places result in vHt_q ndarray. 

38 """ 

39 assert dens.rhot_q is not None 

40 epot = self._solve(vHt_q, dens.rhot_q) 

41 return epot 

42 

43 def _solve(self, 

44 vHt_q: Array1D, 

45 rhot_q: Array1D) -> float: 

46 vHt_q[:] = 4 * pi * self.strength * rhot_q 

47 if self.pd.gd.comm.rank == 0: 

48 # Use uniform backgroud charge in case we have a charged system: 

49 vHt_q[0] = 0.0 

50 vHt_q /= self.G2_q 

51 epot = 0.5 * self.pd.integrate(vHt_q, rhot_q) 

52 return epot 

53 

54 

55class ChargedReciprocalSpacePoissonSolver(ReciprocalSpacePoissonSolver): 

56 def __init__(self, 

57 pd: PWDescriptor, 

58 charge: float, 

59 alpha: float = None, 

60 eps: float = 1e-5): 

61 """Reciprocal-space Poisson solver for charged molecules. 

62 

63 * Add a compensating Guassian-shaped charge to the density 

64 in order to make the total charge neutral (placed in the 

65 middle of the unit cell 

66 

67 * Solve Poisson equation. 

68 

69 * Correct potential so that it has the correct 1/r 

70 asymptotic behavior 

71 

72 * Correct energy to remove the artificial interaction with 

73 the compensation charge 

74 """ 

75 ReciprocalSpacePoissonSolver.__init__(self, pd, charge) 

76 self.charge = charge 

77 

78 if alpha is None: 

79 # Shortest distance from center to edge of cell: 

80 rcut = 0.5 / (pd.gd.icell_cv**2).sum(axis=1).max()**0.5 

81 

82 # Make sure e^(-alpha*rcut^2)=eps: 

83 alpha = -rcut**-2 * np.log(eps) 

84 

85 self.alpha = alpha 

86 

87 center_v = pd.gd.cell_cv.sum(axis=0) / 2 

88 G2_q = pd.G2_qG[0] 

89 G_qv = pd.get_reciprocal_vectors() 

90 self.charge_q = np.exp(-1 / (4 * alpha) * G2_q + 

91 1j * (G_qv @ center_v)) 

92 self.charge_q *= charge / pd.gd.dv 

93 

94 R_Rv = pd.gd.get_grid_point_coordinates().transpose((1, 2, 3, 0)) 

95 d_R = ((R_Rv - center_v)**2).sum(axis=3)**0.5 

96 

97 with seterr(invalid='ignore'): 

98 potential_R = erf(alpha**0.5 * d_R) / d_R 

99 if ((pd.gd.N_c % 2) == 0).all(): 

100 R_c = pd.gd.N_c // 2 

101 if pd.gd.is_my_grid_point(R_c): 

102 potential_R[tuple(R_c - pd.gd.beg_c)] = (4 * alpha / pi)**0.5 

103 self.potential_q = charge * pd.fft(potential_R) 

104 

105 def __str__(self): 

106 return ('Using Gaussian-shaped compensation charge: e^(-ar^2) ' 

107 f'with a={self.alpha:.3f} bohr^-2') 

108 

109 def _solve(self, 

110 vHt_q: Array1D, 

111 rhot_q: Array1D) -> float: 

112 neutral_q = rhot_q + self.charge_q 

113 if self.pd.gd.comm.rank == 0: 

114 error = neutral_q[0] * self.pd.gd.dv 

115 assert error.imag == 0.0, error 

116 assert abs(error.real) < 0.01, error 

117 neutral_q[0] = 0.0 

118 

119 vHt_q[:] = 4 * pi * neutral_q 

120 vHt_q /= self.G2_q 

121 epot = 0.5 * self.pd.integrate(vHt_q, neutral_q) 

122 epot -= self.pd.integrate(self.potential_q, rhot_q) 

123 epot -= self.charge**2 * (self.alpha / 2 / pi)**0.5 

124 vHt_q -= self.potential_q 

125 return epot