Coverage for gpaw/pes/continuum.py: 35%

1from math import pi, sin, cos 

2import numpy as np 

3 

4from ase.units import Bohr 

5import gpaw.cgpaw as cgpaw 

6from gpaw.poisson import PoissonSolver 

7from gpaw.pes.state import State 

8from gpaw.analyse.expandyl import AngularIntegral 

9 

10 

11class PlaneWave(State): 

12 

13 """Plane wave state.""" 

14 

15 def get_grid(self, k_c, r0=None): 

16 if r0 is None: 

17 r0_c = np.array([0., 0., 0.]) 

18 else: 

19 r0_c = np.array(r0) 

20 pw_G = self.gd.zeros(dtype=complex) 

21 python = True 

22 python = False 

23 if python: 

24 r_c = np.array([0., 0., 0.]) 

25 gd = self.gd 

26 for i in range(gd.n_c[0]): 

27 r_c[0] = (gd.beg_c[0] + i) * gd.h_c[0] 

28 for j in range(gd.n_c[1]): 

29 r_c[1] = (gd.beg_c[1] + j) * gd.h_c[1] 

30 for k in range(gd.n_c[2]): 

31 r_c[2] = (gd.beg_c[2] + k) * gd.h_c[2] 

32 phase = np.dot(k_c, r_c - r0_c) 

33 pw_G[i, j, k] = cos(phase) + 1j * sin(phase) 

34 else: 

35 gd = self.gd 

36 cgpaw.plane_wave_grid(gd.beg_c, gd.end_c, 

37 gd.h_cv.diagonal().copy(), 

38 k_c, r0_c, pw_G) 

39 pw_G /= (2 * pi) ** (3. / 2.) 

40 return pw_G 

41 

42 

43class ZerothOrder1(State): 

44 

45 def __init__(self, calculator): 

46 self.calculator = calculator 

47 State.__init__(self, calculator.gd) 

48 self.pw = PlaneWave(self.gd) 

49 

50 # get effective potential 

51 hamiltonian = self.calculator.hamiltonian 

52 if 1: 

53 self.vt_G = hamiltonian.vt_sG[0] # XXX treat spin 

54 else: 

55 self.vt_G = np.where(hamiltonian.vt_sG[0] > 0, 

56 0.0, hamiltonian.vt_sG[0]) 

57 self.intvt = self.gd.integrate(self.vt_G) 

58 print('# int(vt_G)=', self.intvt, (self.vt_G > 0).any()) 

59 

60 self.solve() 

61 

62 def get_grid(self, k_c, r0): 

63 print('Correction:', self.gd.integrate(self.corrt_G)) 

64 if not hasattr(self, 'written'): 

65 print('r0', r0, r0 * Bohr) 

66 dR = 0.3 

67 AI = AngularIntegral(r0 * Bohr, self.gd, dR=dR) 

68 r_R = AI.radii() 

69 psi_R = AI.average(self.corrt_G) 

70 v_R = AI.average(self.vt_G) 

71 with open('radial_dR' + str(dR) + '.dat', 'w') as fd: 

72 print('# R v(R) psi(R)', file=fd) 

73 for r, psi, v in zip(r_R, psi_R, v_R): 

74 print(r, psi, v, file=fd) 

75 

76 self.written = True 

77 

78 return self.pw.get_grid(k_c, r0) - 1e6 * self.corrt_G 

79 

80 def solve(self): 

81 hamiltonian = self.calculator.hamiltonian 

82 self.poisson = PoissonSolver('fd', nn=hamiltonian.poisson.nn) 

83 self.poisson.set_grid_descriptor(self.gd) 

84 self.poisson.initialize() 

85 

86 corrt_G = self.gd.empty() 

87 self.poisson.solve(corrt_G, self.vt_G, charge=None) 

88 corrt_G /= (2 * pi) ** (5. / 2) 

89 

90 self.corrt_G = corrt_G