Coverage for gpaw/analyse/simple_stm.py: 88%

1from math import sqrt 

2 

3import numpy as np 

4from ase.atoms import Atoms 

5from ase.units import Bohr, Hartree 

6from ase.io.cube import read_cube_data, write_cube 

7from ase.dft.stm import STM 

8 

9import gpaw.mpi as mpi 

10from gpaw.grid_descriptor import GridDescriptor 

11 

12 

13class SimpleStm(STM): 

14 

15 """Simple STM object to simulate STM pictures. 

16 

17 The simulation uses either a single pseudo-wavefunction (PWF) 

18 or the PWFs inside the given bias range.""" 

19 

20 def __init__(self, atoms): 

21 

22 self.filename = None 

23 self.is_wf = False 

24 self.bias = None 

25 self.ldos = None 

26 self.heights = None 

27 

28 if isinstance(atoms, str): 

29 self.read_3D(atoms) 

30 self.calc = None 

31 self.atoms = None 

32 else: 

33 if isinstance(atoms, Atoms): 

34 self.atoms = atoms 

35 self.calc = atoms.calc 

36 else: 

37 self.calc = atoms 

38 self.atoms = self.calc.get_atoms() 

39 self.calc.converge_wave_functions() 

40 

41 self.gd = self.calc.wfs.gd 

42 self.offset_c = [int(not a) for a in self.gd.pbc_c] 

43 

44 super().__init__(self.atoms) 

45 

46 def calculate_ldos(self, bias): 

47 """bias is the n, k, s list/tuple.""" 

48 if self.calc is None: 

49 return 

50 

51 self.bias = bias 

52 self.is_wf = True 

53 self.ldos = self.gd.zeros() 

54 

55 if hasattr(bias, '__len__') and len(bias) == 3: 

56 n, k, s = bias 

57 # only a single wf requested 

58 kd = self.calc.wfs.kd 

59 rank, q = kd.who_has(k) 

60 if kd.comm.rank == rank: 

61 u = q * kd.nspins + s 

62 self.add_wf_to_ldos(n, u, weight=1) 

63 

64 else: 

65 # energy bias 

66 try: 

67 efermi_s = self.calc.get_fermi_levels() 

68 except ValueError: 

69 efermi_s = np.array([self.calc.get_fermi_level()] * 2) 

70 

71 if isinstance(bias, (int, float)): 

72 # bias given 

73 if bias > 0: 

74 # positive bias = negative tip 

75 # -> probe unoccupied states 

76 emin_s = efermi_s 

77 emax_s = efermi_s + bias 

78 occupied = False 

79 else: 

80 # negative bias = positive tip 

81 # -> probe occupied states 

82 emin_s = efermi_s + bias 

83 emax_s = efermi_s 

84 occupied = True 

85 else: 

86 # emin and emax given 

87 emin, emax = bias 

88 if abs(emin) > abs(emax): 

89 occupied = True 

90 else: 

91 occupied = False 

92 emin_s = np.array([emin + efermi_s] * 2) 

93 emax_s = np.array([emax + efermi_s] * 2) 

94 

95 emin_s /= Hartree 

96 emax_s /= Hartree 

97 

98 for u in range(len(self.calc.wfs.kpt_u)): 

99 kpt = self.calc.wfs.kpt_u[u] 

100 emin = emin_s[kpt.s] 

101 emax = emax_s[kpt.s] 

102 for n, eps in enumerate(kpt.eps_n): 

103 if eps > emin and eps < emax: 

104 if occupied: 

105 weight = kpt.f_n[n] 

106 else: 

107 weight = kpt.weight - kpt.f_n[n] 

108 self.add_wf_to_ldos(n, u, weight) 

109 

110 def add_wf_to_ldos(self, n, u, weight=None): 

111 """Add the wf with given kpoint and spin to the ldos""" 

112 kpt = self.calc.wfs.kpt_u[u] 

113 psi = kpt.psit_nG[n] 

114 w = weight 

115 if w is None: 

116 w = kpt.weight 

117# print "w=", w, kpt.weight 

118 self.ldos += w * (psi * np.conj(psi)).real 

119 

120 def write_3D(self, bias, filename, filetype=None): 

121 """Write the density as a 3D file. 

122 

123 Units: [e/A^3]""" 

124 self.calculate_ldos(bias) 

125 self.calc.wfs.kd.comm.sum(self.ldos) 

126 ldos = self.gd.collect(self.ldos) 

127# print "write: integrated =", self.gd.integrate(self.ldos) 

128 

129 if mpi.rank != 0: 

130 return 

131 

132 if filetype is None: 

133 # estimate file type from name ending 

134 filetype = filename.split('.')[-1] 

135 filetype = filetype.lower() 

136 

137 if filetype == 'cube': 

138 with open(filename, 'w') as fd: 

139 write_cube(fd, self.calc.get_atoms(), ldos / Bohr ** 3) 

140 else: 

141 raise NotImplementedError('unknown file type "' + filetype + '"') 

142 

143 def read_3D(self, filename, filetype=None): 

144 """Read the density from a 3D file""" 

145 

146 if filetype is None: 

147 # estimate file type from name ending 

148 filetype = filename.split('.')[-1] 

149 

150 if filetype == 'cube': 

151 data, atoms = read_cube_data(filename) 

152 self.cell = atoms.cell.array 

153 

154 pbc_c = [True, True, True] 

155 N_c = np.array(data.shape) 

156 for c in range(3): 

157 if N_c[c] % 2 == 1: 

158 pbc_c[c] = False 

159 N_c[c] += 1 

160 self.gd = GridDescriptor(N_c, self.cell.diagonal() / Bohr, pbc_c) 

161 self.offset_c = [int(not a) for a in self.gd.pbc_c] 

162 

163 else: 

164 raise NotImplementedError('unknown file type "' + filetype + '"') 

165 

166 self.filename = filename 

167 self.ldos = np.array(data * Bohr ** 3) 

168# print "read: integrated =", self.gd.integrate(self.ldos) 

169 

170 def current_to_density(self, current): 

171 """The connection between density n and current I 

172 

173 n [e/Angstrom^3] = 0.0002 sqrt(I [nA]) 

174 

175 as given in Hofer et al., RevModPhys 75 (2003) 1287 

176 """ 

177 return 0.0002 * sqrt(current) 

178 

179 def linescan(self, bias, current, p1, p2, npoints=50, z0=None): 

180 """Line scan for a given current [nA]""" 

181 return STM.linescan(self, bias, 

182 self.current_to_density(current), 

183 p1, p2, npoints, z0) 

184 

185 def density_to_current(self, density): 

186 return 5000. * density ** 2 

187 

188 def scan_const_current(self, current, bias): 

189 """Get the height image for constant current I [nA]. 

190 """ 

191 return self.scan_const_density(self.current_to_density(current), 

192 bias) 

193 

194 def scan_const_density(self, density, bias): 

195 """Get the height image for constant density [e/Angstrom^3]. 

196 """ 

197 x, y, heights = self.scan(bias, density) 

198 self.heights = heights / Bohr 

199 

200 return self.heights