Coverage for gpaw/test/lrtddft/test_excited_state.py: 87%

1import time 

2import pytest 

3import numpy as np 

4 

5from ase import Atom, Atoms, io 

6from ase.parallel import parprint, paropen 

7from ase.units import Ha 

8 

9from gpaw import GPAW 

10from gpaw.mpi import world 

11from gpaw.lrtddft import LrTDDFT 

12from gpaw.lrtddft.excited_state import ExcitedState 

13 

14 

15def get_H2(calculator=None): 

16 """Define H2 and set calculator if given""" 

17 R = 0.7 # approx. experimental bond length 

18 a = 3.0 

19 c = 4.0 

20 H2 = Atoms([Atom('H', (a / 2, a / 2, (c - R) / 2)), 

21 Atom('H', (a / 2, a / 2, (c + R) / 2))], 

22 cell=(a, a, c)) 

23 

24 if calculator is not None: 

25 H2.calc = calculator 

26 

27 return H2 

28 

29 

30def get_H3(calculator=None): 

31 R = 0.87 # approx. bond length 

32 a = 4.0 

33 c = 3.0 

34 H3 = Atoms('H3', positions=[[0, 0, 0], [R, 0, 0], 

35 [R / 2, R / 2 * np.sqrt(3), 0]], 

36 cell=(a, a, c)) 

37 H3.center() 

38 

39 if calculator is not None: 

40 H3.calc = calculator 

41 

42 return H3 

43 

44 

45@pytest.mark.lrtddft 

46def test_split(in_tmp_dir): 

47 fname = 'exlst.out' 

48 calc = GPAW(mode='fd', xc='PBE', h=0.25, nbands=3, txt=fname) 

49 exlst = LrTDDFT(calc, txt=fname) 

50 exst = ExcitedState(exlst, 0, txt=fname) 

51 H2 = get_H2(exst) 

52 H2.get_potential_energy() 

53 

54 n = world.size 

55 exst.split(n) 

56 H2.get_potential_energy() 

57 

58 if world.rank == 0: 

59 with open(fname) as f: 

60 string = f.read() 

61 assert 'Total number of cores used: {0}'.format(n) in string 

62 assert 'Total number of cores used: 1' in string 

63 

64 

65@pytest.mark.lrtddft 

66def test_lrtddft_excited_state(): 

67 txt = None 

68 

69 calc = GPAW(mode='fd', xc='PBE', h=0.25, nbands=3, spinpol=False, txt=txt) 

70 H2 = get_H2(calc) 

71 

72 xc = 'LDA' 

73 lr = LrTDDFT(calc, xc=xc) 

74 

75 # excited state with forces 

76 accuracy = 0.015 

77 exst = ExcitedState(lr, 0, d=0.01, 

78 parallel=2) 

79 

80 t0 = time.time() 

81 parprint("########### first call to forces --> calculate") 

82 forces = exst.get_forces(H2) 

83 parprint("time used:", time.time() - t0) 

84 for c in range(2): 

85 assert forces[0, c] == pytest.approx(0.0, abs=accuracy) 

86 assert forces[1, c] == pytest.approx(0.0, abs=accuracy) 

87 assert forces[0, 2] + forces[1, 2] == pytest.approx(0.0, abs=accuracy) 

88 

89 parprint("########### second call to potential energy --> just return") 

90 t0 = time.time() 

91 E = exst.get_potential_energy() 

92 parprint("E=", E) 

93 parprint("time used:", time.time() - t0) 

94 t0 = time.time() 

95 E = exst.get_potential_energy(H2) 

96 parprint("E=", E) 

97 parprint("time used:", time.time() - t0) 

98 

99 parprint("########### second call to forces --> just return") 

100 t0 = time.time() 

101 exst.get_forces() 

102 parprint("time used:", time.time() - t0) 

103 t0 = time.time() 

104 exst.get_forces(H2) 

105 parprint("time used:", time.time() - t0) 

106 

107 parprint("########### moved atoms, call to forces --> calculate") 

108 p = H2.get_positions() 

109 p[1, 1] += 0.1 

110 H2.set_positions(p) 

111 

112 t0 = time.time() 

113 exst.get_forces(H2) 

114 parprint("time used:", time.time() - t0) 

115 

116 

117@pytest.mark.lrtddft 

118def test_io(in_tmp_dir): 

119 """Test output and input from files""" 

120 calc = GPAW(mode='fd', xc='PBE', h=0.25, nbands=3, txt=None) 

121 exlst = LrTDDFT(calc, txt=None) 

122 exst = ExcitedState(exlst, 0, txt=None) 

123 H2 = get_H2(exst) 

124 

125 parprint('----------- calculate') 

126 E1 = H2.get_potential_energy() 

127 E0 = exst.calculator.get_potential_energy() 

128 dE1 = exlst[0].energy * Ha 

129 assert E1 == pytest.approx(E0 + dE1, 1.e-5) 

130 

131 parprint('----------- write trajectory') 

132 ftraj = 'H2exst.traj' 

133 F = H2.get_forces() 

134 traj = io.Trajectory(ftraj, 'w') 

135 traj.write(H2) 

136 

137 parprint('----------- write') 

138 fname = 'exst_test_io' 

139 parprint('----', exst.get_potential_energy()) 

140 exst.write(fname) 

141 world.barrier() 

142 

143 parprint('----------- read') 

144 exst = ExcitedState.read(fname, txt=None) 

145 E1 = exst.get_potential_energy() 

146 parprint('-----', exst.get_potential_energy(), E0 + dE1) 

147 assert E1 == pytest.approx(E0 + dE1, 1.e-5) 

148 

149 parprint('----------- read trajectory') 

150 atoms = io.read(ftraj) 

151 assert atoms.get_potential_energy() == pytest.approx(E1, 1.e-5) 

152 assert atoms.get_forces() == pytest.approx(F, 1.e-5) 

153 

154 

155@pytest.mark.lrtddft 

156def test_log(in_tmp_dir): 

157 fname = 'ex0_silent.out' 

158 calc = GPAW(mode='fd', xc='PBE', h=0.25, nbands=5, txt=None) 

159 calc.calculate(get_H2(calc)) 

160 exlst = LrTDDFT(calc, restrict={'eps': 0.4, 'jend': 3}, txt=None) 

161 exst = ExcitedState(exlst, 0, txt=fname) 

162 del calc 

163 del exlst 

164 del exst 

165 world.barrier() 

166 

167 if world.rank == 0: 

168 with open(fname) as f: 

169 string = f.read() 

170 assert 'ExcitedState' in string 

171 assert ' ___ ___ ___ _ _ _' not in string 

172 assert 'Linear response TDDFT calculation' not in string 

173 

174 fname = 'ex0_split.out' 

175 calc = GPAW(mode='fd', xc='PBE', h=0.25, nbands=5, txt=fname) 

176 calc.calculate(get_H2(calc)) 

177 exlst = LrTDDFT(calc, restrict={'eps': 0.4, 'jend': 3}, log=calc.log) 

178 exst = ExcitedState(exlst, 0, log=exlst.log, parallel=2) 

179 exst.get_forces() 

180 del calc 

181 del exlst 

182 del exst 

183 

184 if world.rank == 0: 

185 with paropen(fname) as f: 

186 string = f.read() 

187 assert 'ExcitedState' in string 

188 if world.size == 1: 

189 # one eq + 6 * 2 displacements + one eq. = 14 calculations 

190 n = 14 

191 else: 

192 # we see only half of the calculations in parallel 

193 # one eq + 3 * 2 displacements + one eq. = 8 calculations 

194 n = 8 

195 assert string.count('Converged after') == n 

196 assert string.count('Kohn-Sham single transitions') == n 

197 assert string.count('Linear response TDDFT calculation') == n 

198 

199 

200@pytest.mark.lrtddft 

201def test_forces(): 

202 """Test whether force calculation works""" 

203 calc = GPAW(mode='fd', xc='PBE', h=0.25, nbands=3, txt=None) 

204 exlst = LrTDDFT(calc) 

205 exst = ExcitedState(exlst, 0) 

206 H2 = get_H2(exst) 

207 

208 parprint('---------------- serial') 

209 

210 forces = H2.get_forces() 

211 accuracy = 1.e-3 

212 # forces in x and y direction should be 0 

213 assert forces[:, :2] == pytest.approx( 

214 np.zeros_like(forces[:, :2]), abs=accuracy) 

215 

216 # forces in z direction should be opposite 

217 assert -forces[0, 2] == pytest.approx(forces[1, 2], abs=accuracy) 

218 

219 # next call should give back the stored forces 

220 forces1 = exst.get_forces(H2) 

221 assert (forces1 == forces).all() 

222 

223 # test parallel 

224 if world.size > 1: 

225 parprint('---------------- parallel', world.size) 

226 exstp = ExcitedState(exlst, 0, parallel=2) 

227 forcesp = exstp.get_forces(H2) 

228 assert forcesp == pytest.approx(forces, abs=0.001) 

229 

230 

231@pytest.mark.lrtddft 

232def test_unequal_parralel_work(): 

233 """Test whether parallel force calculation works for three atoms""" 

234 if world.size == 1: 

235 return 

236 

237 calc = GPAW(mode='fd', xc='PBE', charge=1, h=0.25, nbands=3, txt=None) 

238 exlst = LrTDDFT(calc, txt=None) 

239 H3 = get_H3() 

240 

241 parprint('---------------- serial') 

242 exst = ExcitedState(exlst, 0, txt=None) 

243 forces = exst.get_forces(H3) 

244 

245 parprint('---------------- parallel', world.size) 

246 exst = ExcitedState(exlst, 0, parallel=2, txt=None) 

247 H3 = get_H3(exst) 

248 forcesp = H3.get_forces() 

249 

250 assert forcesp == pytest.approx(forces, abs=0.01)