Coverage for /builds/ase/ase/ase/dft/stm.py: 82.72%

162 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-08-02 00:12 +0000

1# fmt: off 

2 

3import numpy as np 

4 

5from ase.io.jsonio import read_json, write_json 

6 

7 

8class STM: 

9 def __init__(self, atoms, symmetries=None, use_density=False): 

10 """Scanning tunneling microscope. 

11 

12 atoms: Atoms object or filename 

13 Atoms to scan or name of file to read LDOS from. 

14 symmetries: list of int 

15 List of integers 0, 1, and/or 2 indicating which surface 

16 symmetries have been used to reduce the number of k-points 

17 for the DFT calculation. The three integers correspond to 

18 the following three symmetry operations:: 

19 

20 [-1 0] [ 1 0] [ 0 1] 

21 [ 0 1] [ 0 -1] [ 1 0] 

22 

23 use_density: bool 

24 Use the electron density instead of the LDOS. 

25 """ 

26 

27 self.use_density = use_density 

28 

29 if isinstance(atoms, str): 

30 with open(atoms) as fd: 

31 self.ldos, self.bias, self.cell = read_json(fd, 

32 always_array=False) 

33 self.atoms = None 

34 else: 

35 self.atoms = atoms 

36 self.cell = atoms.cell 

37 self.bias = None 

38 self.ldos = None 

39 assert not self.cell[2, :2].any() and not self.cell[:2, 2].any() 

40 

41 self.symmetries = symmetries or [] 

42 

43 def calculate_ldos(self, bias): 

44 """Calculate local density of states for given bias.""" 

45 if self.ldos is not None and bias == self.bias: 

46 return 

47 

48 self.bias = bias 

49 

50 calc = self.atoms.calc 

51 

52 if self.use_density: 

53 self.ldos = calc.get_pseudo_density() 

54 return 

55 

56 if bias < 0: 

57 emin = bias 

58 emax = 0.0 

59 else: 

60 emin = 0 

61 emax = bias 

62 

63 nbands = calc.get_number_of_bands() 

64 weights = calc.get_k_point_weights() 

65 nkpts = len(weights) 

66 nspins = calc.get_number_of_spins() 

67 eigs = np.array([[calc.get_eigenvalues(k, s) 

68 for k in range(nkpts)] 

69 for s in range(nspins)]) 

70 eigs -= calc.get_fermi_level() 

71 ldos = np.zeros(calc.get_pseudo_wave_function(0, 0, 0).shape) 

72 

73 for s in range(nspins): 

74 for k in range(nkpts): 

75 for n in range(nbands): 

76 e = eigs[s, k, n] 

77 if emin < e < emax: 

78 psi = calc.get_pseudo_wave_function(n, k, s) 

79 ldos += weights[k] * (psi * np.conj(psi)).real 

80 

81 if 0 in self.symmetries: 

82 # (x,y) -> (-x,y) 

83 ldos[1:] += ldos[:0:-1].copy() 

84 ldos[1:] *= 0.5 

85 

86 if 1 in self.symmetries: 

87 # (x,y) -> (x,-y) 

88 ldos[:, 1:] += ldos[:, :0:-1].copy() 

89 ldos[:, 1:] *= 0.5 

90 

91 if 2 in self.symmetries: 

92 # (x,y) -> (y,x) 

93 ldos += ldos.transpose((1, 0, 2)).copy() 

94 ldos *= 0.5 

95 

96 self.ldos = ldos 

97 

98 def write(self, filename): 

99 """Write local density of states to JSON file.""" 

100 write_json(filename, (self.ldos, self.bias, self.cell)) 

101 

102 def get_averaged_current(self, bias, z): 

103 """Calculate avarage current at height z (in Angstrom). 

104 

105 Use this to get an idea of what current to use when scanning.""" 

106 

107 self.calculate_ldos(bias) 

108 nz = self.ldos.shape[2] 

109 

110 # Find grid point: 

111 n = z / self.cell[2, 2] * nz 

112 dn = n - np.floor(n) 

113 n = int(n) % nz 

114 

115 # Average and do linear interpolation: 

116 return ((1 - dn) * self.ldos[:, :, n].mean() + 

117 dn * self.ldos[:, :, (n + 1) % nz].mean()) 

118 

119 def scan(self, bias, current, z0=None, repeat=(1, 1)): 

120 """Constant current 2-d scan. 

121 

122 Returns three 2-d arrays (x, y, z) containing x-coordinates, 

123 y-coordinates and heights. These three arrays can be passed to 

124 matplotlibs contourf() function like this: 

125 

126 >>> import matplotlib.pyplot as plt 

127 >>> plt.contourf(x, y, z) 

128 >>> plt.show() 

129 

130 """ 

131 

132 self.calculate_ldos(bias) 

133 

134 L = self.cell[2, 2] 

135 nz = self.ldos.shape[2] 

136 h = L / nz 

137 

138 ldos = self.ldos.reshape((-1, nz)) 

139 

140 heights = np.empty(ldos.shape[0]) 

141 for i, a in enumerate(ldos): 

142 heights[i] = find_height(a, current, h, z0) 

143 

144 s0 = heights.shape = self.ldos.shape[:2] 

145 heights = np.tile(heights, repeat) 

146 s = heights.shape 

147 

148 ij = np.indices(s, dtype=float).reshape((2, -1)).T 

149 x, y = np.dot(ij / s0, self.cell[:2, :2]).T.reshape((2,) + s) 

150 

151 return x, y, heights 

152 

153 def scan2(self, bias, z, repeat=(1, 1)): 

154 """Constant height 2-d scan. 

155 

156 Returns three 2-d arrays (x, y, I) containing x-coordinates, 

157 y-coordinates and currents. These three arrays can be passed to 

158 matplotlibs contourf() function like this: 

159 

160 >>> import matplotlib.pyplot as plt 

161 >>> plt.contourf(x, y, I) 

162 >>> plt.show() 

163 

164 """ 

165 

166 self.calculate_ldos(bias) 

167 

168 nz = self.ldos.shape[2] 

169 ldos = self.ldos.reshape((-1, nz)) 

170 

171 current = np.empty(ldos.shape[0]) 

172 

173 zp = z / self.cell[2, 2] * nz 

174 zp = int(zp) % nz 

175 

176 for i, a in enumerate(ldos): 

177 current[i] = self.find_current(a, zp) 

178 

179 s0 = current.shape = self.ldos.shape[:2] 

180 current = np.tile(current, repeat) 

181 s = current.shape 

182 

183 ij = np.indices(s, dtype=float).reshape((2, -1)).T 

184 x, y = np.dot(ij / s0, self.cell[:2, :2]).T.reshape((2,) + s) 

185 

186 # Returing scan with axes in Angstrom. 

187 return x, y, current 

188 

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

190 """Constant current line scan. 

191 

192 Example:: 

193 

194 stm = STM(...) 

195 z = ... # tip position 

196 c = stm.get_averaged_current(-1.0, z) 

197 stm.linescan(-1.0, c, (1.2, 0.0), (1.2, 3.0)) 

198 """ 

199 

200 heights = self.scan(bias, current, z0)[2] 

201 

202 p1 = np.asarray(p1, float) 

203 p2 = np.asarray(p2, float) 

204 d = p2 - p1 

205 s = np.dot(d, d)**0.5 

206 

207 cell = self.cell[:2, :2] 

208 shape = np.array(heights.shape, float) 

209 M = np.linalg.inv(cell) 

210 line = np.empty(npoints) 

211 for i in range(npoints): 

212 p = p1 + i * d / (npoints - 1) 

213 q = np.dot(p, M) * shape 

214 line[i] = interpolate(q, heights) 

215 return np.linspace(0, s, npoints), line 

216 

217 def pointcurrent(self, bias, x, y, z): 

218 """Current for a single x, y, z position for a given bias.""" 

219 

220 self.calculate_ldos(bias) 

221 

222 nx = self.ldos.shape[0] 

223 ny = self.ldos.shape[1] 

224 nz = self.ldos.shape[2] 

225 

226 # Find grid point: 

227 xp = x / np.linalg.norm(self.cell[0]) * nx 

228 dx = xp - np.floor(xp) 

229 xp = int(xp) % nx 

230 

231 yp = y / np.linalg.norm(self.cell[1]) * ny 

232 dy = yp - np.floor(yp) 

233 yp = int(yp) % ny 

234 

235 zp = z / np.linalg.norm(self.cell[2]) * nz 

236 dz = zp - np.floor(zp) 

237 zp = int(zp) % nz 

238 

239 # 3D interpolation of the LDOS at point (x,y,z) at given bias. 

240 xyzldos = (((1 - dx) + (1 - dy) + (1 - dz)) * self.ldos[xp, yp, zp] + 

241 dx * self.ldos[(xp + 1) % nx, yp, zp] + 

242 dy * self.ldos[xp, (yp + 1) % ny, zp] + 

243 dz * self.ldos[xp, yp, (zp + 1) % nz]) 

244 

245 return dos2current(bias, xyzldos) 

246 

247 def sts(self, x, y, z, bias0, bias1, biasstep): 

248 """Returns the dI/dV curve for position x, y at height z (in Angstrom), 

249 for bias from bias0 to bias1 with step biasstep.""" 

250 

251 biases = np.arange(bias0, bias1 + biasstep, biasstep) 

252 current = np.zeros(biases.shape) 

253 

254 for b in np.arange(len(biases)): 

255 print(b, biases[b]) 

256 current[b] = self.pointcurrent(biases[b], x, y, z) 

257 

258 dIdV = np.gradient(current, biasstep) 

259 

260 return biases, current, dIdV 

261 

262 def find_current(self, ldos, z): 

263 """ Finds current for given LDOS at height z.""" 

264 nz = self.ldos.shape[2] 

265 

266 zp = z / self.cell[2, 2] * nz 

267 dz = zp - np.floor(zp) 

268 zp = int(zp) % nz 

269 

270 ldosz = (1 - dz) * ldos[zp] + dz * ldos[(zp + 1) % nz] 

271 

272 return dos2current(self.bias, ldosz) 

273 

274 

275def dos2current(bias, dos): 

276 # Borrowed from gpaw/analyse/simple_stm.py: 

277 # The connection between density n and current I 

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

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

280 return 5000. * dos**2 * (1 if bias > 0 else -1) 

281 

282 

283def interpolate(q, heights): 

284 qi = q.astype(int) 

285 f = q - qi 

286 g = 1 - f 

287 qi %= heights.shape 

288 n0, m0 = qi 

289 n1, m1 = (qi + 1) % heights.shape 

290 z = (g[0] * g[1] * heights[n0, m0] + 

291 f[0] * g[1] * heights[n1, m0] + 

292 g[0] * f[1] * heights[n0, m1] + 

293 f[0] * f[1] * heights[n1, m1]) 

294 return z 

295 

296 

297def find_height(ldos, current, h, z0=None): 

298 if z0 is None: 

299 n = len(ldos) - 2 

300 else: 

301 n = int(z0 / h) 

302 while n >= 0: 

303 if ldos[n] > current: 

304 break 

305 n -= 1 

306 else: 

307 return 0.0 

308 

309 c2, c1 = ldos[n:n + 2] 

310 return (n + 1 - (current - c1) / (c2 - c1)) * h 

311 

312 

313def delta(biases, bias, width): 

314 """Return a delta-function centered at 'bias'""" 

315 x = -((biases - bias) / width)**2 

316 return np.exp(x) / (np.sqrt(np.pi) * width)