Coverage for gpaw/new/fd/builder.py: 97%

1from __future__ import annotations 

2 

3from gpaw.core import UGArray, UGDesc 

4from gpaw.new.builder import create_uniform_grid 

5from gpaw.new.fd.hamiltonian import FDHamiltonian 

6from gpaw.new.fd.pot_calc import FDPotentialCalculator 

7from gpaw.new.gpw import as_double_precision 

8from gpaw.new.pwfd.builder import PWFDDFTComponentsBuilder 

9 

10 

11class FDDFTComponentsBuilder(PWFDDFTComponentsBuilder): 

12 def __init__(self, 

13 atoms, 

14 params, 

15 *, 

16 comm, 

17 log): 

18 self.interpolation_stencil_range = params.interpolation or 3 

19 super().__init__(atoms, 

20 params, 

21 comm=comm, log=log) 

22 assert not self.soc 

23 assert self.qspiral_v is None 

24 

25 self._nct_aR = None 

26 self._tauct_aR = None 

27 

28 self.electrostatic_potential_desc = self.fine_grid 

29 self.interpolation_desc = self.fine_grid 

30 

31 def create_uniform_grids(self): 

32 grid = create_uniform_grid( 

33 self.mode, 

34 self.params.gpts, 

35 self.atoms.cell, 

36 self.atoms.pbc, 

37 self.ibz.symmetries, 

38 h=self.params.h, 

39 interpolation=self.interpolation_stencil_range, 

40 comm=self.communicators['d']) 

41 fine_grid = grid.new(size=grid.size_c * 2) 

42 # decomposition=[2 * d for d in grid.decomposition] 

43 return grid, fine_grid 

44 

45 def create_wf_description(self) -> UGDesc: 

46 return self.grid.new(dtype=self.dtype) 

47 

48 def get_pseudo_core_densities(self): 

49 if self._nct_aR is None: 

50 self._nct_aR = self.setups.create_pseudo_core_densities( 

51 self.grid, self.relpos_ac, atomdist=self.atomdist, xp=self.xp) 

52 return self._nct_aR 

53 

54 def get_pseudo_core_ked(self): 

55 if self._tauct_aR is None: 

56 self._tauct_aR = self.setups.create_pseudo_core_ked( 

57 self.grid, self.relpos_ac, atomdist=self.atomdist) 

58 return self._tauct_aR 

59 

60 def create_potential_calculator(self): 

61 env = self.create_environment(self.fine_grid) 

62 poisson_solver = env.create_poisson_solver( 

63 grid=self.fine_grid, xp=self.xp, 

64 solver=self.params.poissonsolver) 

65 return FDPotentialCalculator( 

66 self.grid, self.fine_grid, self.setups, self.xc, poisson_solver, 

67 relpos_ac=self.relpos_ac, atomdist=self.atomdist, 

68 interpolation_stencil_range=self.interpolation_stencil_range, 

69 environment=env, 

70 extensions=self.get_extensions(), 

71 xp=self.xp) 

72 

73 def create_hamiltonian_operator(self): 

74 return FDHamiltonian(self.wf_desc, 

75 kin_stencil=self.params.mode.nn, 

76 xp=self.xp) 

77 

78 def convert_wave_functions_from_uniform_grid(self, 

79 C_nM, 

80 basis_set, 

81 kpt_c, 

82 q): 

83 grid = self.grid.new(kpt=kpt_c, dtype=self.dtype) 

84 psit_nR = grid.zeros(self.nbands, self.communicators['b']) 

85 mynbands = len(C_nM.data) 

86 basis_set.lcao_to_grid(C_nM.data, psit_nR.data[:mynbands], q) 

87 return psit_nR.to_xp(self.xp) 

88 

89 def read_ibz_wave_functions(self, reader): 

90 ibzwfs = super().read_ibz_wave_functions(reader) 

91 

92 if 'coefficients' in reader.wave_functions: 

93 name = 'coefficients' 

94 elif 'values' in reader.wave_functions: 

95 name = 'values' # old name 

96 else: 

97 return ibzwfs 

98 

99 singlep = reader.get('precision', 'double') == 'single' 

100 c = reader.bohr**1.5 

101 if reader.version < 0: 

102 c = 1 # old gpw file 

103 

104 for wfs in ibzwfs: 

105 grid = self.wf_desc.new(kpt=wfs.kpt_c) 

106 index = (wfs.spin, wfs.k) 

107 data = reader.wave_functions.proxy(name, *index) 

108 data.scale = c 

109 if self.communicators['w'].size == 1 and not singlep: 

110 wfs.psit_nX = UGArray(grid, self.nbands, data=data) 

111 else: 

112 band_comm = self.communicators['b'] 

113 wfs.psit_nX = UGArray( 

114 grid, self.nbands, 

115 comm=band_comm) 

116 if grid.comm.rank == 0: 

117 mynbands = (self.nbands + 

118 band_comm.size - 1) // band_comm.size 

119 n1 = min(band_comm.rank * mynbands, self.nbands) 

120 n2 = min((band_comm.rank + 1) * mynbands, self.nbands) 

121 assert wfs.psit_nX.mydims[0] == n2 - n1 

122 data = data[n1:n2] # read from file 

123 

124 if singlep: 

125 wfs.psit_nX.scatter_from(as_double_precision(data)) 

126 else: 

127 wfs.psit_nX.scatter_from(data) 

128 

129 return ibzwfs