1import numpy as np 

2 

3from ase.calculators.calculator import Calculator, all_properties 

4from ase.calculators.calculator import PropertyNotImplementedError 

5 

6 

7class SinglePointCalculator(Calculator): 

8 """Special calculator for a single configuration. 

9 

10 Used to remember the energy, force and stress for a given 

11 configuration. If the positions, atomic numbers, unit cell, or 

12 boundary conditions are changed, then asking for 

13 energy/forces/stress will raise an exception.""" 

14 

15 name = 'unknown' 

16 

17 def __init__(self, atoms, **results): 

18 """Save energy, forces, stress, ... for the current configuration.""" 

19 Calculator.__init__(self) 

20 self.results = {} 

21 for property, value in results.items(): 

22 assert property in all_properties 

23 if value is None: 

24 continue 

25 if property in ['energy', 'magmom', 'free_energy']: 

26 self.results[property] = value 

27 else: 

28 self.results[property] = np.array(value, float) 

29 self.atoms = atoms.copy() 

30 

31 def __str__(self): 

32 tokens = [] 

33 for key, val in sorted(self.results.items()): 

34 if np.isscalar(val): 

35 txt = '{}={}'.format(key, val) 

36 else: 

37 txt = '{}=...'.format(key) 

38 tokens.append(txt) 

39 return '{}({})'.format(self.__class__.__name__, ', '.join(tokens)) 

40 

41 def get_property(self, name, atoms=None, allow_calculation=True): 

42 if atoms is None: 

43 atoms = self.atoms 

44 if name not in self.results or self.check_state(atoms): 

45 if allow_calculation: 

46 raise PropertyNotImplementedError( 

47 'The property "{0}" is not available.'.format(name)) 

48 return None 

49 

50 result = self.results[name] 

51 if isinstance(result, np.ndarray): 

52 result = result.copy() 

53 return result 

54 

55 

56class SinglePointKPoint: 

57 def __init__(self, weight, s, k, eps_n=[], f_n=[]): 

58 self.weight = weight 

59 self.s = s # spin index 

60 self.k = k # k-point index 

61 self.eps_n = eps_n 

62 self.f_n = f_n 

63 

64 

65def arrays_to_kpoints(eigenvalues, occupations, weights): 

66 """Helper function for building SinglePointKPoints. 

67 

68 Convert eigenvalue, occupation, and weight arrays to list of 

69 SinglePointKPoint objects.""" 

70 nspins, nkpts, nbands = eigenvalues.shape 

71 assert eigenvalues.shape == occupations.shape 

72 assert len(weights) == nkpts 

73 kpts = [] 

74 for s in range(nspins): 

75 for k in range(nkpts): 

76 kpt = SinglePointKPoint( 

77 weight=weights[k], s=s, k=k, 

78 eps_n=eigenvalues[s, k], f_n=occupations[s, k]) 

79 kpts.append(kpt) 

80 return kpts 

81 

82 

83class SinglePointDFTCalculator(SinglePointCalculator): 

84 def __init__(self, atoms, 

85 efermi=None, bzkpts=None, ibzkpts=None, bz2ibz=None, 

86 **results): 

87 self.bz_kpts = bzkpts 

88 self.ibz_kpts = ibzkpts 

89 self.bz2ibz = bz2ibz 

90 self.eFermi = efermi 

91 

92 SinglePointCalculator.__init__(self, atoms, **results) 

93 self.kpts = None 

94 

95 def get_fermi_level(self): 

96 """Return the Fermi-level(s).""" 

97 return self.eFermi 

98 

99 def get_bz_to_ibz_map(self): 

100 return self.bz2ibz 

101 

102 def get_bz_k_points(self): 

103 """Return the k-points.""" 

104 return self.bz_kpts 

105 

106 def get_number_of_spins(self): 

107 """Return the number of spins in the calculation. 

108 

109 Spin-paired calculations: 1, spin-polarized calculation: 2.""" 

110 if self.kpts is not None: 

111 nspin = set() 

112 for kpt in self.kpts: 

113 nspin.add(kpt.s) 

114 return len(nspin) 

115 return None 

116 

117 def get_spin_polarized(self): 

118 """Is it a spin-polarized calculation?""" 

119 nos = self.get_number_of_spins() 

120 if nos is not None: 

121 return nos == 2 

122 return None 

123 

124 def get_ibz_k_points(self): 

125 """Return k-points in the irreducible part of the Brillouin zone.""" 

126 return self.ibz_kpts 

127 

128 def get_kpt(self, kpt=0, spin=0): 

129 if self.kpts is not None: 

130 counter = 0 

131 for kpoint in self.kpts: 

132 if kpoint.s == spin: 

133 if kpt == counter: 

134 return kpoint 

135 counter += 1 

136 return None 

137 

138 def get_k_point_weights(self): 

139 """ Retunrs the weights of the k points """ 

140 if self.kpts is not None: 

141 weights = [] 

142 for kpoint in self.kpts: 

143 if kpoint.s == 0: 

144 weights.append(kpoint.weight) 

145 return np.array(weights) 

146 return None 

147 

148 def get_occupation_numbers(self, kpt=0, spin=0): 

149 """Return occupation number array.""" 

150 kpoint = self.get_kpt(kpt, spin) 

151 if kpoint is not None: 

152 return kpoint.f_n 

153 return None 

154 

155 def get_eigenvalues(self, kpt=0, spin=0): 

156 """Return eigenvalue array.""" 

157 kpoint = self.get_kpt(kpt, spin) 

158 if kpoint is not None: 

159 return kpoint.eps_n 

160 return None 

161 

162 def get_homo_lumo(self): 

163 """Return HOMO and LUMO energies.""" 

164 if self.kpts is None: 

165 raise RuntimeError('No kpts') 

166 eHs = [] 

167 eLs = [] 

168 for kpt in self.kpts: 

169 eH, eL = self.get_homo_lumo_by_spin(kpt.s) 

170 eHs.append(eH) 

171 eLs.append(eL) 

172 return np.array(eHs).max(), np.array(eLs).min() 

173 

174 def get_homo_lumo_by_spin(self, spin=0): 

175 """Return HOMO and LUMO energies for a given spin.""" 

176 if self.kpts is None: 

177 raise RuntimeError('No kpts') 

178 for kpt in self.kpts: 

179 if kpt.s == spin: 

180 break 

181 else: 

182 raise RuntimeError('No k-point with spin {0}'.format(spin)) 

183 if self.eFermi is None: 

184 raise RuntimeError('Fermi level is not available') 

185 eH = -1.e32 

186 eL = 1.e32 

187 for kpt in self.kpts: 

188 if kpt.s == spin: 

189 for e in kpt.eps_n: 

190 if e <= self.eFermi: 

191 eH = max(eH, e) 

192 else: 

193 eL = min(eL, e) 

194 return eH, eL