Coverage for gpaw/inducedfield/inducedfield_lrtddft.py: 85%

72 statements  

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

1import numpy as np 

2from copy import copy 

3 

4from ase.parallel import parprint 

5 

6from gpaw.lrtddft import LrTDDFT 

7from gpaw.utilities.folder import Folder 

8 

9from gpaw.inducedfield.inducedfield_base import BaseInducedField 

10 

11 

12class LrTDDFTInducedField(BaseInducedField): 

13 """Induced field class for linear response TDDFT. 

14 

15 Attributes (see also ``BaseInducedField``): 

16 ------------------------------------------- 

17 lr: LrTDDFT object 

18 LrTDDFT object 

19 """ 

20 

21 def __init__(self, filename=None, paw=None, lr=None, 

22 frequencies=None, folding='Gauss', width=0.08, 

23 kickdir=0 

24 ): 

25 """ 

26 Parameters (see also ``BaseInducedField``): 

27 ------------------------------------------- 

28 paw: GPAW object 

29 GPAW object for ground state 

30 lr: LrTDDFT object 

31 LrTDDFT object 

32 kickdir: int 

33 Kick direction: 0, 1, 2 for x, y, z 

34 """ 

35 # TODO: change kickdir to general kick = [1e-3, 1-e3, 0] etc. 

36 

37 if not isinstance(lr, LrTDDFT): 

38 raise RuntimeError('Provide LrTDDFT object.') 

39 

40 # Check that lr is diagonalized 

41 if len(lr) == 0: 

42 raise RuntimeError('Diagonalize LrTDDFT first.') 

43 

44 self.lr = lr 

45 

46 # "Kick" direction 

47 self.kickdir = kickdir 

48 

49 BaseInducedField.__init__(self, filename, paw, 

50 frequencies, folding, width) 

51 

52 self.readwritemode_str_to_list = \ 

53 {'': ['Frho', 'atoms'], 

54 'all': ['Frho', 'Fphi', 'Fef', 'Ffe', 'atoms'], 

55 'field': ['Frho', 'Fphi', 'Fef', 'Ffe', 'atoms']} 

56 

57 def initialize(self, paw, allocate=True): 

58 BaseInducedField.initialize(self, paw, allocate) 

59 

60 # Artificial background electric field (linear response) 

61 # TODO: change kickdir to general kick = [1e-3, 1-e3, 0] etc. 

62 Fbgef_v = np.zeros((self.nv,), dtype=float) 

63 Fbgef_v[self.kickdir] = 1.0 

64 self.Fbgef_v = Fbgef_v 

65 

66 # Initialize PAW object 

67 paw.initialize_positions() 

68 

69 def set_folding(self, folding, width): 

70 BaseInducedField.set_folding(self, folding, width) 

71 

72 if self.folding is None: 

73 self.folder = Folder(None, None) 

74 else: 

75 if self.folding == 'Gauss': 

76 self.folder = Folder(self.width, 'ComplexGauss') 

77 elif self.folding == 'Lorentz': 

78 self.folder = Folder(self.width, 'ComplexLorentz') 

79 else: 

80 raise RuntimeError('unknown folding "' + self.folding + '"') 

81 

82 def get_induced_density(self, from_density='comp', gridrefinement=2): 

83 

84 paw = self.paw 

85 lr = self.lr 

86 omega_w = self.omega_w 

87 

88 if gridrefinement == 1: 

89 gd = self.gd 

90 elif gridrefinement == 2: 

91 gd = self.density.finegd 

92 else: 

93 raise NotImplementedError 

94 

95 nkss = len(lr.kss) 

96 nexcs = len(lr) 

97 omega_I = np.empty((nexcs), dtype=float) 

98 

99 # C_Iij: weights for each excitation I and KS-pair ij 

100 C_Iij = np.zeros((nexcs, nkss), dtype=float) 

101 

102 # Calculate C_Iij 

103 FIsqfe_ij = np.zeros((nkss), dtype=float) 

104 for I, exc in enumerate(lr): 

105 omega_I[I] = exc.energy 

106 

107 B = 0.0 

108 for ij, ks in enumerate(lr.kss): 

109 FIsqfe_ij[ij] = exc.f[ij] * np.sqrt(ks.fij * ks.energy) 

110 B += ks.mur[self.kickdir] * FIsqfe_ij[ij] 

111 

112 C_Iij[I] = 2 * B * FIsqfe_ij 

113 

114 # Fold C_Iij to C_wij 

115 # C_wij: weights for each requested frequency w and KS-pair ij 

116 self.omega_w, C_wij = self.folder.fold_values(omega_I, C_Iij, omega_w) 

117 

118 assert (self.omega_w == omega_w).all() # TODO: remove 

119 

120 Fn_wg = gd.zeros((self.nw), dtype=complex) 

121 

122 kpt_u = paw.wfs.kpt_u 

123 

124 for ij, ks_ in enumerate(lr.kss): 

125 # TODO: better output of progress 

126 parprint('%d/%d' % (ij, nkss)) 

127 

128 # TODO: speedup: 

129 # for each w, check magnitude of C_wij and 

130 # take C_wij for those ij that contribute most 

131 

132 # Take copy so that memory for pair density is released 

133 # after each loop iteration 

134 ks = copy(ks_) 

135 

136 # Initialize pair density 

137 kpt = kpt_u[ks.spin] 

138 ks.set_paw(paw) 

139 ks.initialize(kpt, ks.i, ks.j) 

140 

141 # Get pair density 

142 # TODO: gridrefinements 

143 yes_finegrid = gridrefinement == 2 

144 if from_density == 'pseudo': 

145 nij_g = ks.get(finegrid=yes_finegrid) 

146 elif from_density == 'comp': 

147 nij_g = ks.with_compensation_charges(finegrid=yes_finegrid) 

148 elif from_density == 'ae': 

149 nij_g = ks.with_ae_corrections(finegrid=yes_finegrid) 

150 # TODO: better to split pair density in pseudo density part 

151 # and FD_asp etc. coefficients (like in time-propagation). 

152 # Then do gridrefinement and add AE/comp corrections afterwards 

153 # -> speedup (no need to do summations on fine grid) 

154 

155 # TODO: speedup: check magnitude of C_wij (see above) 

156 for w in range(self.nw): 

157 Fn_wg[w] += nij_g * C_wij[w, ij] 

158 

159 # Return charge density (electrons = negative charge) 

160 return - Fn_wg, gd