Coverage for gpaw/test/solvation/test_poisson.py: 100%

113 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-07-09 00:21 +0000

1import warnings 

2 

3from ase.parallel import parprint 

4from ase.units import Bohr 

5import numpy as np 

6from scipy.special import erf 

7 

8from gpaw.solvation.poisson import (WeightedFDPoissonSolver, 

9 ADM12PoissonSolver, 

10 PolarizationPoissonSolver) 

11from gpaw.solvation.dielectric import Dielectric 

12from gpaw.grid_descriptor import GridDescriptor 

13from gpaw.utilities.gauss import Gaussian 

14from gpaw.fd_operators import Gradient 

15import pytest 

16 

17nn = 3 

18accuracy = 2e-10 

19 

20 

21def gradient(gd, x, nn): 

22 out = gd.empty(3) 

23 for i in (0, 1, 2): 

24 Gradient(gd, i, 1.0, nn).apply(x, out[i]) 

25 return out 

26 

27 

28def make_gd(h, box, pbc): 

29 diag = np.array([box] * 3) 

30 cell = np.diag(diag) 

31 grid_shape = tuple((diag / h * 2).astype(int)) 

32 return GridDescriptor(grid_shape, cell / Bohr, pbc) 

33 

34 

35class MockDielectric(Dielectric): 

36 def __init__(self, epsinf, nn): 

37 Dielectric.__init__(self, epsinf) 

38 self.nn = nn 

39 

40 def update(self, cavity): 

41 grad_eps_vg = gradient(self.gd, self.eps_gradeps[0] - self.epsinf, 

42 self.nn) 

43 for i in (0, 1, 2): 

44 self.eps_gradeps[1 + i][...] = grad_eps_vg[i] 

45 

46 

47def test_solvation_poisson(): 

48 box = 12. 

49 gd = make_gd(h=.4, box=box, pbc=False) 

50 

51 def solve(ps, eps, rho): 

52 phi = gd.zeros() 

53 dielectric = MockDielectric(epsinf=eps.max(), nn=nn) 

54 dielectric.set_grid_descriptor(gd) 

55 dielectric.allocate() 

56 dielectric.eps_gradeps[0][...] = eps 

57 dielectric.update(None) 

58 with warnings.catch_warnings(): 

59 # ignore production code warning for alternative psolvers 

60 warnings.simplefilter("ignore") 

61 solver = ps(nn=nn, relax='J', eps=accuracy) 

62 solver.set_dielectric(dielectric) 

63 solver.set_grid_descriptor(gd) 

64 solver.solve(phi, rho) 

65 return phi 

66 

67 psolvers = (WeightedFDPoissonSolver, 

68 ADM12PoissonSolver, 

69 PolarizationPoissonSolver) 

70 

71 # test neutral system with constant permittivity 

72 parprint('neutral, constant permittivity') 

73 epsinf = 80. 

74 eps = gd.zeros() 

75 eps.fill(epsinf) 

76 qs = (-1., 1.) 

77 shifts = (-1., 1.) 

78 rho = gd.zeros() 

79 phi_expected = gd.zeros() 

80 for q, shift in zip(qs, shifts): 

81 gauss_norm = q / np.sqrt(4 * np.pi) 

82 gauss = Gaussian(gd, center=(box / 2. + shift) * np.ones(3) / Bohr) 

83 rho += gauss_norm * gauss.get_gauss(0) 

84 phi_expected += gauss_norm * gauss.get_gauss_pot(0) / epsinf 

85 

86 for ps in psolvers: 

87 phi = solve(ps, eps, rho) 

88 parprint(ps, np.abs(phi - phi_expected).max()) 

89 assert phi == pytest.approx(phi_expected, abs=1e-3) 

90 

91 # test charged system with constant permittivity 

92 parprint('charged, constant permittivity') 

93 epsinf = 80. 

94 eps = gd.zeros() 

95 eps.fill(epsinf) 

96 q = -2. 

97 gauss_norm = q / np.sqrt(4 * np.pi) 

98 gauss = Gaussian(gd, center=(box / 2. + 1.) * np.ones(3) / Bohr) 

99 rho_gauss = gauss_norm * gauss.get_gauss(0) 

100 phi_gauss = gauss_norm * gauss.get_gauss_pot(0) 

101 phi_expected = phi_gauss / epsinf 

102 

103 for ps in psolvers: 

104 phi = solve(ps, eps, rho_gauss) 

105 parprint(ps, np.abs(phi - phi_expected).max()) 

106 assert phi == pytest.approx(phi_expected, abs=1e-3) 

107 

108 # test non-constant permittivity 

109 msgs = ('neutral, non-constant permittivity', 

110 'charged, non-constant permittivity') 

111 qss = ((-1., 1.), (2., )) 

112 shiftss = ((-.4, .4), (-1., )) 

113 epsshifts = (.0, -1.) 

114 

115 for msg, qs, shifts, epsshift in zip(msgs, qss, shiftss, epsshifts): 

116 parprint(msg) 

117 epsinf = 80. 

118 gauss = Gaussian(gd, center=(box / 2. + epsshift) * np.ones(3) / Bohr) 

119 eps = gauss.get_gauss(0) 

120 eps = epsinf - eps / eps.max() * (epsinf - 1.) 

121 

122 rho = gd.zeros() 

123 phi_expected = gd.zeros() 

124 grad_eps = gradient(gd, eps - epsinf, nn) 

125 

126 for q, shift in zip(qs, shifts): 

127 gauss = Gaussian(gd, center=(box / 2. + shift) * np.ones(3) / Bohr) 

128 phi_tmp = gauss.get_gauss_pot(0) 

129 xyz = gauss.xyz 

130 fac = 2. * np.sqrt(gauss.a) * np.exp(-gauss.a * gauss.r2) 

131 fac /= np.sqrt(np.pi) * gauss.r2 

132 fac -= erf(np.sqrt(gauss.a) * gauss.r) / (gauss.r2 * gauss.r) 

133 fac *= 2.0 * 1.7724538509055159 

134 grad_phi = fac * xyz 

135 laplace_phi = -4. * np.pi * gauss.get_gauss(0) 

136 rho_tmp = -1. / (4. * np.pi) * ( 

137 (grad_eps * grad_phi).sum(0) + eps * laplace_phi) 

138 norm = gd.integrate(rho_tmp) 

139 rho_tmp /= norm * q 

140 phi_tmp /= norm * q 

141 rho += rho_tmp 

142 phi_expected += phi_tmp 

143 

144 for ps in psolvers: 

145 phi = solve(ps, eps, rho) 

146 parprint(ps, np.abs(phi - phi_expected).max()) 

147 assert phi == pytest.approx(phi_expected, abs=1e-3)