Coverage for gpaw/helmholtz.py: 89%

1from math import pi, sqrt 

2import numpy as np 

3from numpy import exp 

4from scipy.special import erf, erfc 

5 

6from gpaw.poisson import FDPoissonSolver 

7from gpaw.utilities.gauss import Gaussian 

8from gpaw.fd_operators import FDOperator, laplace 

9from gpaw.transformers import Transformer 

10 

11 

12class HelmholtzGaussian(Gaussian): 

13 def get_phi(self, k2): 

14 """Get the solution of the Helmholtz equation for a Gaussian.""" 

15# This should lead to very big errors 

16 r = np.sqrt(self.r2) 

17 k = sqrt(k2) 

18 sigma = 1. / sqrt(2 * self.a) 

19 p = sigma * k / sqrt(2) 

20 i = 1j 

21 rhop = r / sqrt(2) / sigma + i * p 

22 rhom = r / sqrt(2) / sigma - i * p 

23 # h = np.sin(k * r) 

24 

25 # the gpaw-Gaussian is sqrt(4 * pi) times a 3D normalized Gaussian 

26 return sqrt(4 * pi) * exp(-p**2) / r / 2 * ( 

27 np.cos(k * r) * (erf(rhop) + erf(rhom)) + 

28 i * np.sin(k * r) * (2 + erf(rhop) - erf(rhom))) 

29 

30 

31class ScreenedPoissonGaussian(Gaussian): 

32 def get_phi(self, mu2): 

33 """Get the solution of the screened Poisson equation for a Gaussian. 

34 

35 The Gaussian is centered to middle of grid-descriptor.""" 

36 r = np.sqrt(self.r2) 

37 mu = sqrt(mu2) 

38 sigma = 1. / sqrt(2 * self.a) 

39 sig2 = sigma**2 

40 mrho = (sig2 * mu - r) / (sqrt(2) * sigma) 

41 prho = (sig2 * mu + r) / (sqrt(2) * sigma) 

42 

43 # the gpaw-Gaussian is sqrt(4 * pi) times a 3D normalized Gaussian 

44 return sqrt(4 * pi) * exp(sig2 * mu2 / 2.0) / (2 * r) * ( 

45 exp(-mu * r) * erfc(mrho) - exp(mu * r) * erfc(prho)) 

46 

47 

48class HelmholtzOperator(FDOperator): 

49 def __init__(self, gd, scale=1.0, n=1, dtype=float, k2=0.0): 

50 """Helmholtz for general non orthorhombic grid. 

51 

52 gd: GridDescriptor 

53 Descriptor for grid. 

54 scale: float 

55 Scaling factor. Use scale=-0.5 for a kinetic energy operator. 

56 n: int 

57 Range of stencil. Stencil has O(h^(2n)) error. 

58 dtype: float or complex 

59 Data-type to work on. 

60 """ 

61 

62 # Order the 13 neighbor grid points: 

63 M_ic = np.indices((3, 3, 3)).reshape((3, -3)).T[-13:] - 1 

64 u_cv = gd.h_cv / (gd.h_cv**2).sum(1)[:, np.newaxis]**0.5 

65 u2_i = (np.dot(M_ic, u_cv)**2).sum(1) 

66 i_d = u2_i.argsort() 

67 

68 m_mv = np.array([(2, 0, 0), (0, 2, 0), (0, 0, 2), 

69 (0, 1, 1), (1, 0, 1), (1, 1, 0)]) 

70 # Try 3, 4, 5 and 6 directions: 

71 for D in range(3, 7): 

72 h_dv = np.dot(M_ic[i_d[:D]], gd.h_cv) 

73 A_md = (h_dv**m_mv[:, np.newaxis, :]).prod(2) 

74 a_d, residual, rank, s = np.linalg.lstsq(A_md, [1, 1, 1, 0, 0, 0], 

75 rcond=-1) 

76 if residual.sum() < 1e-14: 

77 assert rank == D, 'You have a weird unit cell!' 

78 # D directions was OK 

79 break 

80 

81 a_d *= scale 

82 offsets = [(0, 0, 0)] 

83 coefs = [laplace[n][0] * a_d.sum()] 

84 coefs[0] += k2 * scale 

85 for d in range(D): 

86 M_c = M_ic[i_d[d]] 

87 offsets.extend(np.arange(1, n + 1)[:, np.newaxis] * M_c) 

88 coefs.extend(a_d[d] * np.array(laplace[n][1:])) 

89 offsets.extend(np.arange(-1, -n - 1, -1)[:, np.newaxis] * M_c) 

90 coefs.extend(a_d[d] * np.array(laplace[n][1:])) 

91 

92 FDOperator.__init__(self, coefs, offsets, gd, dtype) 

93 

94 self.description = ( 

95 '%d*%d+1=%d point O(h^%d) finite-difference Helmholtz' % 

96 ((self.npoints - 1) // n, n, self.npoints, 2 * n)) 

97 

98 

99class HelmholtzSolver(FDPoissonSolver): 

100 """Solve the Helmholtz or screened Poisson equations. 

101 

102 The difference between the Helmholtz equation: 

103 

104 (Laplace + k^2) phi = n 

105 

106 and the screened Poisson equation: 

107 

108 (Laplace - mu^2) phi = n 

109 

110 is only the sign of the added inhomogenity. Because of 

111 this we can use one class to solve both. So if k2 is 

112 greater zero we'll try to solve the Helmhlotz equation, 

113 otherwise we'll try to solve the screened Poisson equation. 

114 """ 

115 

116 def __init__(self, k2=0.0, nn='M', relax='GS', eps=2e-10, 

117 use_charge_center=True): 

118 assert k2 <= 0, 'Currently only defined for k^2<=0' 

119 FDPoissonSolver.__init__(self, nn, relax, eps, 

120 use_charge_center=use_charge_center) 

121 self.k2 = k2 

122 

123 def set_grid_descriptor(self, gd): 

124 # Should probably be renamed initialize 

125 self.gd = gd 

126 self.dv = gd.dv 

127 

128 gd = self.gd 

129 scale = -0.25 / pi 

130 

131 if self.nn == 'M': 

132 raise ValueError( 

133 'Helmholtz not defined for Mehrstellen stencil') 

134 self.operators = [HelmholtzOperator(gd, scale, self.nn, k2=self.k2)] 

135 self.B = None 

136 

137 self.interpolators = [] 

138 self.restrictors = [] 

139 

140 level = 0 

141 self.presmooths = [2] 

142 self.postsmooths = [1] 

143 

144 # Weights for the relaxation, 

145 # only used if 'J' (Jacobi) is chosen as method 

146 self.weights = [2.0 / 3.0] 

147 

148 while level < 4: 

149 try: 

150 gd2 = gd.coarsen() 

151 except ValueError: 

152 break 

153 self.operators.append(HelmholtzOperator(gd2, scale, 1, 

154 k2=self.k2)) 

155 self.interpolators.append(Transformer(gd2, gd)) 

156 self.restrictors.append(Transformer(gd, gd2)) 

157 self.presmooths.append(4) 

158 self.postsmooths.append(4) 

159 self.weights.append(1.0) 

160 level += 1 

161 gd = gd2 

162 

163 self.levels = level 

164 

165 if self.relax_method == 1: 

166 self.description = 'Gauss-Seidel' 

167 else: 

168 self.description = 'Jacobi' 

169 self.description += ' solver with %d multi-grid levels' % (level + 1) 

170 self.description += '\nStencil: ' + self.operators[0].description 

171 

172 def load_gauss(self, center=None): 

173 """Load the gaussians.""" 

174 if not hasattr(self, 'rho_gauss') or center is not None: 

175 if self.k2 > 0: 

176 gauss = HelmholtzGaussian(self.gd, center=center) 

177 else: 

178 gauss = ScreenedPoissonGaussian(self.gd, center=center) 

179 self.rho_gauss = gauss.get_gauss(0) 

180 self.phi_gauss = gauss.get_phi(abs(self.k2))