Coverage for gpaw/test/poisson/test_bloechl.py: 88%

1from math import pi 

2 

3import numpy as np 

4import pytest 

5from ase import Atoms 

6from scipy.special import erf 

7 

8from gpaw.atom.radialgd import EquidistantRadialGridDescriptor as RGD 

9from gpaw.core import PWDesc 

10from gpaw.new.ase_interface import GPAW 

11from gpaw.new.pw.bloechl_poisson import BloechlPAWPoissonSolver 

12from gpaw.new.pw.paw_poisson import (SimplePAWPoissonSolver, 

13 SlowPAWPoissonSolver) 

14from gpaw.new.pw.poisson import PWPoissonSolver 

15from gpaw.mpi import world 

16from gpaw.gpu import cupy as cp 

17from gpaw.gpu.mpi import CuPyMPI 

18 

19 

20def g(rc, rgd): 

21 """Gaussian.""" 

22 return rgd.spline(4 / rc**3 / np.pi**0.5 * np.exp(-(rgd.r_g / rc)**2), 

23 l=0) 

24 

25 

26def c(r, rc1, rc2): 

27 """Coulomb interaction between 2 gaussians.""" 

28 a1 = 1 / rc1**2 

29 a2 = 1 / rc2**2 

30 f = 2 * (pi**5 / (a1 + a2))**0.5 / (a1 * a2) 

31 f *= 16 / pi / rc1**3 / rc2**3 

32 if r == 0.0: 

33 return f 

34 T = a1 * a2 / (a1 + a2) * r**2 

35 y = 0.5 * f * erf(T**0.5) * (pi / T)**0.5 

36 return y 

37 

38 

39def energy(charges): 

40 e0 = 0.0 

41 for q1, r1, p1 in charges: 

42 for q2, r2, p2 in charges: 

43 d = abs(p1 - p2) 

44 e12 = 0.5 * q1 * q2 * c(d, r1, r2) / (4 * np.pi)**2 

45 # print(q1, q2, rc1, rc2, d, e12) 

46 e0 += e12 

47 return e0 

48 

49 

50def force(charges, a): 

51 eps = 1e-5 

52 charges[a + 2, 2] += eps 

53 ep = energy(charges) 

54 charges[a + 2, 2] -= 2 * eps 

55 em = energy(charges) 

56 charges[a + 2, 2] += eps 

57 return (em - ep) / (2 * eps) 

58 

59 

60@pytest.mark.parametrize('xp', 

61 [np, 

62 pytest.param(cp, marks=pytest.mark.gpu)]) 

63def test_psolve(xp): 

64 """Unit-test for Blöchl's fast Poisson-solver.""" 

65 comm = CuPyMPI(world) 

66 rgd = RGD(0.01, 500) 

67 rc1 = 0.6 

68 rc2 = 0.7 

69 d12 = 1.35 

70 g_ai = [[g(rc1, rgd)], [g(rc2, rgd)]] 

71 v = 7.5 

72 gcut = 25.0 

73 pw = PWDesc(gcut=gcut, cell=[2 * v, 2 * v, 2 * v + d12], comm=comm) 

74 relpos_ac = np.array([[0.5, 0.5, v / (2 * v + d12)], 

75 [0.5, 0.5, (v + d12) / (2 * v + d12)]]) 

76 g_aig = pw.atom_centered_functions(g_ai, positions=relpos_ac, xp=xp) 

77 nt_g = pw.zeros(xp=xp) 

78 C_ai = g_aig.empty() 

79 if 0 in C_ai: 

80 C_ai[0] = 0.9 

81 if 1 in C_ai: 

82 C_ai[1] = 0.7 

83 C_ai.data *= 1.0 / (4.0 * np.pi)**0.5 

84 g_aig.add_to(nt_g, C_ai) 

85 

86 charges = np.array( 

87 [(0.9, rc1, 0.0), 

88 (0.7, rc2, d12), 

89 (-0.9, 0.3, 0.0), 

90 (-0.7, 0.4, d12)]) 

91 e0 = energy(charges) 

92 f0 = force(charges, 0) 

93 f1 = force(charges, 1) 

94 print(e0, f0, f1) 

95 

96 ps = PWPoissonSolver(pw) 

97 spps = SimplePAWPoissonSolver( 

98 pw, [0.3, 0.4], ps, relpos_ac, g_aig.atomdist, xp=xp) 

99 Q_aL = spps.ghat_aLg.empty() 

100 Q_aL.data[:] = 0.0 

101 for a, C_i in C_ai.items(): 

102 Q_aL[a][0] = -C_i[0] 

103 nt0_g = nt_g.gather() 

104 vt10_g = pw.zeros(xp=xp).gather() 

105 e1, vHt_g, V1_aL = spps.solve(nt0_g, Q_aL, vt10_g) 

106 F1_av = spps.force_contribution(Q_aL, vHt_g, nt_g) 

107 comm.sum(F1_av) 

108 assert e1 == pytest.approx(e0, abs=1e-9) 

109 print('simple', e1, e1 - e0) 

110 print(F1_av) 

111 assert xp.allclose(F1_av, [[0, 0, f0], 

112 [0, 0, f1]]) 

113 

114 pps = BloechlPAWPoissonSolver( 

115 pw, [0.3, 0.4], ps, relpos_ac, g_aig.atomdist, xp=xp) 

116 vt20_g = pw.zeros(xp=xp).gather() 

117 e2, vHt_g, V2_aL = pps.solve(nt0_g, Q_aL, vt20_g) 

118 F2_av = pps.force_contribution(Q_aL, vHt_g, nt_g) 

119 comm.sum(F2_av) 

120 assert e2 == pytest.approx(e0, abs=1e-8) 

121 print('\nfast ', e2, e2 - e0) 

122 assert xp.allclose(V2_aL.data[::9], V1_aL.data[::9]) 

123 if comm.rank == 0: 

124 vt10_g = vt10_g.to_xp(np) 

125 vt20_g = vt20_g.to_xp(np) 

126 assert vt20_g.data[:5] == pytest.approx(vt10_g.data[:5], abs=1e-10) 

127 assert xp.allclose(F1_av, F2_av, atol=3e-6) 

128 

129 if 0: 

130 ps = PWPoissonSolver(pw.new(gcut=2 * gcut)) 

131 opps = SlowPAWPoissonSolver( 

132 pw, [0.3, 0.4], ps, relpos_ac, g_aig.atomdist) 

133 vt_g = pw.zeros() 

134 e3, vHt_h, V_aL = opps.solve(nt_g, Q_aL, vt_g) 

135 print('old ', e3, e3 - e0) 

136 print(V_aL.data[::9]) 

137 print(vt_g.data[:5]) 

138 

139 

140def fast_slow(fast): 

141 atoms = Atoms('H2', [[0, 0, 0], [0.1, 0.2, 0.8]], pbc=True) 

142 atoms.center(vacuum=3.5) 

143 atoms.calc = GPAW(mode={'name': 'pw', 'ecut': 600}, 

144 poissonsolver={'fast': fast}, 

145 convergence={'forces': 1e-3}, 

146 # txt=None, 

147 symmetry='off') 

148 atoms.get_potential_energy() 

149 f = atoms.get_forces() 

150 eps = 0.001 / 2 

151 atoms.positions[1, 2] += eps 

152 ep = atoms.get_potential_energy() 

153 atoms.positions[1, 2] -= 2 * eps 

154 em = atoms.get_potential_energy() 

155 print(f[1, 2], (em - ep) / (2 * eps)) 

156 

157 

158if __name__ == '__main__': 

159 # test_psolve() 

160 import sys 

161 fast_slow(int(sys.argv[1]))