Coverage for gpaw/test/poisson/test_sjm.py: 97%

1from types import SimpleNamespace 

2import numpy as np 

3from gpaw.core import PWDesc, UGDesc 

4from gpaw.new.sjm import SJMPWPoissonSolver 

5from gpaw.new.pw.poisson import ConjugateGradientPoissonSolver 

6 

7 

8def f(a, z, z0, w): 

9 return np.exp(-((z - z0) / w)**2) / a**2 / np.pi**0.5 / w 

10 

11 

12if 0: # Analytic result 

13 from sympy import Symbol, exp, integrate, oo, var 

14 z = var('z') 

15 w = Symbol('w', positive=True) 

16 m = integrate(exp(-(z / w)**2), (z, -oo, oo)) 

17 print(m) # sqrt(pi)*w 

18 

19 

20def test_sjm(): 

21 a = 1.0 

22 L = 10.0 

23 grid = UGDesc.from_cell_and_grid_spacing(cell=[a, a, L], grid_spacing=0.15) 

24 z = grid.xyz()[0, 0, :, 2] 

25 c = 0.05 

26 rhot = f(a, z, 4.0, 1.0) * (1 + c) # electrons 

27 rhot -= f(a, z, 4.0, 0.5) # nucleui 

28 rhot -= f(a, z, 6.0, 1.0) * c # jellium 

29 eps = 1.0 + f(a, z, 5.0, 2.0) * 20 # dielectric 

30 rhot_r = grid.zeros() 

31 rhot_r.data[:] = rhot 

32 eps_r = grid.zeros() 

33 eps_r.data[:] = eps 

34 print(rhot_r.integrate()) 

35 pw = PWDesc(ecut=grid.ekin_max(), cell=grid.cell) 

36 ps = SJMPWPoissonSolver(pw, dielectric=None) 

37 vt_g = pw.zeros() 

38 rhot_g = rhot_r.fft(pw=pw) 

39 ps.solve(vt_g, rhot_g) 

40 vt_r = vt_g.ifft(grid=grid) 

41 dielectric = SimpleNamespace(eps_gradeps=[eps_r.data]) 

42 ps2 = ConjugateGradientPoissonSolver(pw, grid, dielectric, 

43 zero_vacuum=True) 

44 vt2_g = pw.zeros() 

45 ps2.solve(vt2_g, rhot_g) 

46 vt2_r = vt2_g.ifft(grid=grid) 

47 if 0: 

48 import matplotlib.pyplot as plt 

49 plt.plot(z, rhot_r.data[0, 0]) 

50 plt.plot(z, vt_r.data[0, 0]) 

51 plt.plot(z, vt2_r.data[0, 0]) 

52 plt.show() 

53 

54 

55if __name__ == '__main__': 

56 test_sjm()