Coverage for eminus/orbitals.py: 98.59%

1# SPDX-FileCopyrightText: 2021 The eminus developers 

2# SPDX-License-Identifier: Apache-2.0 

3"""Workflow functions that combine functions to generate orbitals.""" 

4 

5from .dft import get_psi 

6from .io import write_cube 

7from .localizer import get_FLO, get_FO, get_scdm, get_wannier 

8from .logger import log 

9from .tools import orbital_center 

10 

11 

12def KSO(scf, write_cubes=False, **kwargs): 

13 """Generate Kohn-Sham orbitals and optionally save them as CUBE files. 

14 

15 Reference: Phys. Rev. 140, A1133. 

16 

17 Args: 

18 scf: SCF object. 

19 

20 Keyword Args: 

21 write_cubes: Write orbitals to CUBE files. 

22 **kwargs: Throwaway arguments. 

23 

24 Returns: 

25 Real-space Kohn-Sham orbitals. 

26 """ 

27 atoms = scf.atoms 

28 

29 # Calculate eigenfunctions and transform to real-space 

30 kso = atoms.I(get_psi(scf, scf.W)) 

31 if write_cubes: 

32 cube_writer(atoms, "KSO", kso) 

33 return kso 

34 

35 

36def FO(scf, write_cubes=False, fods=None, guess="wannier"): 

37 """Generate Fermi orbitals and optionally save them as CUBE files. 

38 

39 Reference: J. Chem. Phys. 153, 084104. 

40 

41 Args: 

42 scf: SCF object. 

43 

44 Keyword Args: 

45 write_cubes: Write orbitals to CUBE files. 

46 fods: Fermi-orbital descriptors. 

47 guess: Guess to generate FODs if none are given. Can be "Wannier" (default) or "PyCOM". 

48 

49 Returns: 

50 Real-space Fermi orbitals. 

51 """ 

52 # Lazy import extras 

53 from .extras.fods import get_fods, remove_core_fods 

54 

55 atoms = scf.atoms 

56 

57 # Calculate eigenfunctions 

58 kso = get_psi(scf, scf.W) 

59 if fods is None and guess == "wannier": 

60 wo = WO(scf) 

61 fods = orbital_center(atoms, wo[0]) 

62 if fods is None and guess == "pycom": 

63 fods = get_fods(atoms) 

64 fods = remove_core_fods(atoms, fods) 

65 

66 # The FO functions need orbitals in reciprocal space as input 

67 fo = get_FO(atoms, kso, fods) 

68 if write_cubes: 

69 cube_writer(atoms, "FO", fo) 

70 return fo 

71 

72 

73def FLO(scf, write_cubes=False, fods=None, guess="wannier"): 

74 """Generate Fermi-Loewdin orbitals and optionally save them as CUBE files. 

75 

76 Reference: J. Chem. Phys. 153, 084104. 

77 

78 Args: 

79 scf: SCF object. 

80 

81 Keyword Args: 

82 write_cubes: Write orbitals to CUBE files. 

83 fods: Fermi-orbital descriptors. 

84 guess: Guess to generate FODs if none are given. Can be "Wannier" (default) or "PyCOM". 

85 

86 Returns: 

87 Real-space Fermi-Loewdin orbitals. 

88 """ 

89 # Lazy import extras 

90 from .extras.fods import get_fods, remove_core_fods 

91 

92 atoms = scf.atoms 

93 

94 # Calculate eigenfunctions 

95 kso = get_psi(scf, scf.W) 

96 

97 guess = guess.lower() 

98 if fods is None and guess == "wannier": 

99 wo = WO(scf) 

100 fods = orbital_center(atoms, wo[0]) 

101 if fods is None and guess == "pycom": 

102 fods = get_fods(atoms) 

103 fods = remove_core_fods(atoms, fods) 

104 

105 # The FLO functions need orbitals in reciprocal space as input 

106 flo = get_FLO(atoms, kso, fods) 

107 if write_cubes: 

108 cube_writer(atoms, "FLO", flo) 

109 return flo 

110 

111 

112def WO(scf, write_cubes=False, precondition=True): 

113 """Generate Wannier orbitals and optionally save them as CUBE files. 

114 

115 Reference: Phys. Rev. B 59, 9703. 

116 

117 Args: 

118 scf: SCF object. 

119 

120 Keyword Args: 

121 write_cubes: Write orbitals to CUBE files. 

122 precondition: Precondition by calculating SCDMs as the initial guess. 

123 

124 Returns: 

125 Real-space Wannier orbitals. 

126 """ 

127 atoms = scf.atoms 

128 

129 # Calculate eigenfunctions/initial guess orbitals and transform to real-space 

130 if precondition: 

131 psi = SCDM(scf) 

132 else: 

133 psi = atoms.I(get_psi(scf, scf.W)) 

134 

135 wo = get_wannier(atoms, psi) 

136 if write_cubes: 

137 cube_writer(atoms, "WO", wo) 

138 return wo 

139 

140 

141def SCDM(scf, write_cubes=False): 

142 """Generate SCDM orbitals and optionally save them as CUBE files. 

143 

144 Reference: J. Chem. Theory Comput. 11, 1463. 

145 

146 Args: 

147 scf: SCF object. 

148 

149 Keyword Args: 

150 write_cubes: Write orbitals to CUBE files. 

151 

152 Returns: 

153 Real-space SCDM orbitals. 

154 """ 

155 atoms = scf.atoms 

156 

157 # Calculate eigenfunctions 

158 kso = get_psi(scf, scf.W) 

159 scdm = get_scdm(atoms, kso) 

160 if write_cubes: 

161 cube_writer(atoms, "SCDM", scdm) 

162 return scdm 

163 

164 

165def cube_writer(atoms, orb_type, orbitals): 

166 """Simple CUBE file writer function. 

167 

168 Args: 

169 atoms: Atoms object. 

170 orb_type: Orbital orb_type for the CUBE file names. 

171 orbitals: Real-space orbitals. 

172 """ 

173 # Create the system name 

174 name = "" 

175 for ia in sorted(set(atoms.atom)): 

176 # Skip the number of atoms if it is equal to one 

177 na = atoms.atom.count(ia) 

178 name += f"{ia}{na if na > 1 else ''}" 

179 

180 n_spin = "" 

181 for ik in range(atoms.kpts.Nk): 

182 for spin in range(atoms.occ.Nspin): 

183 for i in range(atoms.occ.Nstate): 

184 if atoms.occ.f[ik, spin, i] > 0: 

185 if atoms.unrestricted: 

186 n_spin = f"_spin_{spin}" 

187 filename = f"{name}_{orb_type}_k{ik}_{i}{n_spin}.cube" 

188 log.info(f"Write {filename}...") 

189 write_cube(atoms, filename, orbitals[ik][spin, :, i])