Coverage for gpaw/new/pwfd/builder.py: 98%

1from math import pi 

2 

3import numpy as np 

4 

5from gpaw.new.builder import DFTComponentsBuilder 

6from gpaw.new.pwfd.ibzwfs import PWFDIBZWaveFunctions 

7from gpaw.new.lcao.eigensolver import LCAOEigensolver 

8from gpaw.new.lcao.hamiltonian import LCAOHamiltonian 

9from gpaw.new.pwfd.wave_functions import PWFDWaveFunctions 

10from gpaw.core.arrays import XArrayWithNoData 

11 

12 

13class PWFDDFTComponentsBuilder(DFTComponentsBuilder): 

14 def __init__(self, 

15 atoms, 

16 params, 

17 *, 

18 comm=None, 

19 log=None): 

20 super().__init__(atoms, params, comm=comm, log=log) 

21 qspiral = params.mode.qspiral 

22 self.qspiral_v = (None if qspiral is None else 

23 qspiral @ self.grid.icell * (2 * pi)) 

24 

25 def create_eigensolver(self, hamiltonian): 

26 from gpaw.dft import DefaultEigensolver 

27 es = self.params.eigensolver 

28 if isinstance(es, DefaultEigensolver): 

29 es = es.from_param({'name': 'davidson', **es.params}) 

30 return es.build( 

31 self.nbands, 

32 self.wf_desc, 

33 self.communicators['b'], 

34 hamiltonian, 

35 self.params.convergence.get('bands', 'occupied'), 

36 self.setups, 

37 self.atoms) 

38 

39 def read_ibz_wave_functions(self, reader): 

40 kpt_comm, band_comm, domain_comm = (self.communicators[x] 

41 for x in 'kbd') 

42 

43 def create_wfs(spin: int, q: int, k: int, kpt_c, weight: float): 

44 psit_nG = XArrayWithNoData( 

45 comm=band_comm, 

46 dims=(self.nbands,), 

47 desc=self.wf_desc.new(kpt=kpt_c), 

48 xp=self.xp) 

49 wfs = PWFDWaveFunctions( 

50 spin=spin, 

51 q=q, 

52 k=k, 

53 weight=weight, 

54 psit_nX=psit_nG, # type: ignore 

55 setups=self.setups, 

56 relpos_ac=self.relpos_ac, 

57 atomdist=self.atomdist, 

58 ncomponents=self.ncomponents, 

59 qspiral_v=self.qspiral_v) 

60 

61 return wfs 

62 

63 ibzwfs = PWFDIBZWaveFunctions.create( 

64 ibz=self.ibz, 

65 ncomponents=self.ncomponents, 

66 create_wfs_func=create_wfs, 

67 kpt_comm=self.communicators['k'], 

68 kpt_band_comm=self.communicators['D'], 

69 comm=self.communicators['w']) 

70 

71 # Set eigenvalues, occupations, etc.. 

72 self.read_wavefunction_values(reader, ibzwfs) 

73 

74 return ibzwfs 

75 

76 def create_ibz_wave_functions(self, basis, potential): 

77 from gpaw.new.lcao.builder import create_lcao_ibzwfs 

78 

79 if self.params.random: 

80 return self.create_random_ibz_wave_functions() 

81 

82 # sl_default = self.params.parallel['sl_default'] 

83 # sl_lcao = self.params.parallel['sl_lcao'] or sl_default 

84 

85 lcao_dtype = complex if \ 

86 np.issubdtype(self.dtype, np.complexfloating) else float 

87 

88 lcaonbands = min(self.nbands, 

89 basis.Mmax * (2 if self.ncomponents == 4 else 1)) 

90 lcao_ibzwfs, _ = create_lcao_ibzwfs( 

91 basis, 

92 self.ibz, self.communicators, self.setups, 

93 self.relpos_ac, self.grid, lcao_dtype, 

94 lcaonbands, self.ncomponents, self.atomdist, self.nelectrons) 

95 

96 self.log('\nDiagonalizing LCAO Hamiltonian', flush=True) 

97 

98 hamiltonian = LCAOHamiltonian(basis) 

99 LCAOEigensolver(basis).iterate( 

100 lcao_ibzwfs, None, potential, hamiltonian) 

101 

102 self.log('Converting LCAO to grid', flush=True) 

103 

104 def create_wfs(spin, q, k, kpt_c, weight): 

105 lcaowfs = lcao_ibzwfs.wfs_qs[q][spin] 

106 assert lcaowfs.spin == spin 

107 

108 # Convert to PW-coefs in PW-mode: 

109 psit_nX = self.convert_wave_functions_from_uniform_grid( 

110 lcaowfs.C_nM, basis, kpt_c, q) 

111 

112 mylcaonbands, nao = lcaowfs.C_nM.dist.shape 

113 mynbands = len(psit_nX.data) 

114 eig_n = np.empty(self.nbands) 

115 eig_n[:lcaonbands] = lcaowfs._eig_n 

116 eig_n[lcaonbands:] = 100.0 # set high value for random wfs. 

117 if mylcaonbands < mynbands: 

118 psit_nX[mylcaonbands:].randomize( 

119 seed=self.communicators['w'].rank) 

120 wfs = PWFDWaveFunctions( 

121 psit_nX=psit_nX, 

122 spin=spin, 

123 q=q, 

124 k=k, 

125 weight=weight, 

126 setups=self.setups, 

127 relpos_ac=self.relpos_ac, 

128 atomdist=self.atomdist, 

129 ncomponents=self.ncomponents, 

130 qspiral_v=self.qspiral_v) 

131 wfs._eig_n = eig_n 

132 return wfs 

133 

134 return PWFDIBZWaveFunctions.create( 

135 ibz=self.ibz, 

136 ncomponents=self.ncomponents, 

137 create_wfs_func=create_wfs, 

138 kpt_comm=self.communicators['k'], 

139 kpt_band_comm=self.communicators['D'], 

140 comm=self.communicators['w']) 

141 

142 def create_random_ibz_wave_functions(self): 

143 self.log('Initializing wave functions with random numbers') 

144 

145 def create_wfs(spin, q, k, kpt_c, weight): 

146 desc = self.wf_desc.new(kpt=kpt_c) 

147 psit_nX = desc.empty( 

148 dims=(self.nbands,), 

149 comm=self.communicators['b'], 

150 xp=self.xp) 

151 psit_nX.randomize() 

152 

153 wfs = PWFDWaveFunctions( 

154 psit_nX=psit_nX, 

155 spin=spin, 

156 q=q, 

157 k=k, 

158 weight=weight, 

159 setups=self.setups, 

160 relpos_ac=self.relpos_ac, 

161 atomdist=self.atomdist, 

162 ncomponents=self.ncomponents, 

163 qspiral_v=self.qspiral_v) 

164 

165 return wfs 

166 

167 return PWFDIBZWaveFunctions.create( 

168 ibz=self.ibz, 

169 ncomponents=self.ncomponents, 

170 create_wfs_func=create_wfs, 

171 kpt_comm=self.communicators['k'], 

172 kpt_band_comm=self.communicators['D'], 

173 comm=self.communicators['w'])