Coverage for tbkit/lattice.py: 100%

1from __future__ import annotations 

2 

3import numpy as np 

4from numpy.typing import NDArray 

5import matplotlib.pyplot as plt 

6from matplotlib.figure import Figure 

7import tbkit.error_handling as error_handling 

8 

9 

10PI = np.pi 

11COOR_DTYPE = [('x', 'f8'), ('y', 'f8'), ('tag', 'U1')] 

12 

13 

14################################# 

15# CLASS LATTICE 

16################################# 

17 

18 

19class Lattice(): 

20 r''' 

21 Build up 1D or 2D lattice. 

22 Lattice is defined by the discrete operation: 

23 

24 .. math:: 

25 

26 \mathbf{R} = n_1\mathbf{a}_1 + n_2\mathbf{a}_2 

27 

28 where :math:`\mathbf{a}_1` and :math:`\mathbf{a}_2` are the two primitive 

29 vectors and :math:`n_1` and :math:`n_2` are the number of unit cells along 

30 :math:`\mathbf{a}_1` and :math:`\mathbf{a}_2`. 

31 

32 :param unit_cell: List of dictionaries. 

33 One dictionary per site within the unit cell. Each dictionary has two keys: 

34 

35 * 'tag', one-character string. Label of the associated sublattice. 

36 * 'r0', Tuple. Position. 

37 :param prim_vec: List of tuples. 

38 Define the primitive vectors. List of one/two tuples for 1D/2D respectively: 

39 

40 * Tuple, cartesian coordinate of the primitive vector :math:`\mathbf{a}_1`. 

41 * Tuple, cartesian coordinate of the primitive vector :math:`\mathbf{a}_2`. 

42 

43 Example usage:: 

44 

45 # Line-Centered Square lattice 

46 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'a', 'r0': (0., 1.)}] 

47 prim_vec = [(0, 2), (2, 0)] 

48 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec) 

49 ''' 

50 

51 def __init__(self, unit_cell: list[dict], prim_vec: list[tuple[float, float]]) -> None: 

52 error_handling.unit_cell(unit_cell) 

53 error_handling.prim_vec(prim_vec) 

54 self.unit_cell = unit_cell 

55 self.prim_vec = prim_vec 

56 self.tags = np.unique(np.array([dic['tag'] for dic in self.unit_cell])) 

57 self.n1, self.n2 = 0, 0 

58 self.coor = np.array([], dtype=COOR_DTYPE) 

59 self.sites = 0 

60 

61 def get_lattice(self, n1: int, n2: int = 1) -> None: 

62 r''' 

63 Get the lattice positions. 

64 

65 :param n1: Positive Integer. 

66 Number of unit cells along :math:`\mathbf{a}_1`. 

67 :param n2: Positive Integer. Default value 1. 

68 Number of unit cells along :math:`\mathbf{a}_2`. 

69 

70 Example usage:: 

71 

72 # Line-Centered Square lattice 

73 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}, {'tag': 'a', 'r0': (0., 1.)}] 

74 prim_vec = [(0, 2), (2, 0)] 

75 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec) 

76 lat.get_lattice(n1=4, n2=5) 

77 ''' 

78 error_handling.get_lattice(self.prim_vec, n1, n2) 

79 sites_uc = len(self.unit_cell) 

80 sites_tag = n1*n2 

81 self.sites = sites_uc * sites_tag 

82 self.coor = np.empty(self.sites, dtype=COOR_DTYPE) 

83 self.n1, self.n2 = n1, n2 

84 x = self.prim_vec[0][0] * np.arange(n1, dtype='f8') 

85 y = self.prim_vec[0][1] * np.arange(n1, dtype='f8') 

86 xx = np.empty(n1*n2) 

87 yy = np.empty(n1*n2) 

88 xx[:n1] = x 

89 yy[:n1] = y 

90 for i in range(1, n2): 

91 xx[i*n1: (i+1)*n1] = x + i * self.prim_vec[1][0] 

92 yy[i*n1: (i+1)*n1] = y + i * self.prim_vec[1][1] 

93 for i, dic in enumerate(self.unit_cell): 

94 self.coor['x'][i*sites_tag: (i+1)*sites_tag] = xx + dic['r0'][0] 

95 self.coor['y'][i*sites_tag: (i+1)*sites_tag] = yy + dic['r0'][1] 

96 self.coor['tag'][i*sites_tag: (i+1)*sites_tag] = dic['tag'] 

97 self.coor = np.sort(self.coor, order=('y', 'x')) 

98 

99 def add_sites(self, coor: NDArray) -> None: 

100 ''' 

101 Add sites. 

102 

103 :param coor: Structured array with keys: {'x', 'y', 'tag'}. 

104 

105 Example usage:: 

106 

107 # Square lattice 

108 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}] 

109 prim_vec = [(0, 1), (1, 0)] 

110 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec) 

111 lat.get_lattice(n1=2, n2=2) 

112 coor = np.array([(-1., -1, 'b'), (-2., -2, 'c')], 

113 dtype=[('x', 'f8'), ('y', 'f8'), ('tag', 'U1')]) 

114 lat.add_sites(coor) 

115 ''' 

116 error_handling.coor(coor) 

117 self.coor = np.concatenate([self.coor, coor]) 

118 self.sites += len(coor) 

119 self.tags = np.unique(np.concatenate([self.tags, coor['tag']])) 

120 self.coor = np.sort(self.coor, order=('y', 'x')) 

121 

122 def remove_sites(self, index: list[int]) -> None: 

123 ''' 

124 Remove sites defined by their indices 

125 (use method lattice.plot(plt_index=True) 

126 to get access to the site indices). 

127 

128 :param index: List. Site indices to be removed. 

129 

130 Example usage:: 

131 

132 # Square lattice 

133 unit_cell = [{'tag': 'a', 'r0': (0., 0.)}] 

134 prim_vec = [(0, 1), (1, 0)] 

135 lat = lattice(unit_cell=unit_cell, prim_vec=prim_vec) 

136 lat.get_lattice(n1=2, n2=2) 

137 lat.remove_sites([0, 2]) 

138 ''' 

139 error_handling.empty_coor(self.coor) 

140 error_handling.remove_sites(index, self.sites) 

141 mask = np.ones(self.sites, bool) 

142 mask[index] = False 

143 self.coor = self.coor[mask] 

144 self.sites = self.coor.size 

145 

146 def remove_dangling(self) -> None: 

147 ''' 

148 Remove dangling sites 

149 (sites connected with just another site). 

150 ''' 

151 error_handling.empty_coor(self.coor) 

152 while True: 

153 dif_x = self.coor['x'] - self.coor['x'].reshape(self.sites, 1) 

154 dif_y = self.coor['y'] - self.coor['y'].reshape(self.sites, 1) 

155 dis = np.sqrt(dif_x ** 2 + dif_y ** 2) 

156 dis_unique = np.unique(dis) 

157 len_hop = dis_unique[1] 

158 ind = np.argwhere(np.isclose(dis, len_hop)) 

159 dang = [] 

160 for i in range(self.sites): 

161 if (ind[:, 0] == i).sum() == 1: 

162 dang.append(i) 

163 self.coor = np.delete(self.coor, dang, axis=0) 

164 self.sites -= len(dang) 

165 if dang == []: 

166 break 

167 

168 def shift_x(self, shift: float) -> None: 

169 ''' 

170 Shift the x coordinates. 

171 

172 :param shift: Real number. Shift value. 

173 ''' 

174 error_handling.empty_coor(self.coor) 

175 error_handling.real_number(shift, 'shift') 

176 self.coor['x'] += shift 

177 

178 def shift_y(self, shift: float) -> None: 

179 ''' 

180 Shift by *delta_x* the x coordinates. 

181 

182 :param shift: Real number. Shift value. 

183 ''' 

184 error_handling.empty_coor(self.coor) 

185 error_handling.real_number(shift, 'shift') 

186 self.coor['y'] += shift 

187 

188 def change_sign_x(self) -> None: 

189 ''' 

190 Change x coordinates sign. 

191 ''' 

192 error_handling.empty_coor(self.coor) 

193 self.coor['x'] *= -1 

194 

195 def change_sign_y(self) -> None: 

196 ''' 

197 Change y coordinates sign. 

198 ''' 

199 error_handling.empty_coor(self.coor) 

200 self.coor['y'] *= -1 

201 

202 def boundary_line(self, cx: float, cy: float, co: float) -> None: 

203 r''' 

204 Select sites according to :math:`c_yy+c_xx > c_0`. 

205 

206 :param cx: Real number. cx value. 

207 :param cy: Real number. cy value. 

208 :param co: Real number. co value. 

209 ''' 

210 error_handling.empty_coor(self.coor) 

211 error_handling.real_number(cx, 'cx') 

212 error_handling.real_number(cy, 'cy') 

213 error_handling.real_number(co, 'co') 

214 self.coor = self.coor[cy * self.coor['y'] + cx * self.coor['x'] > co] 

215 self.sites = len(self.coor) 

216 

217 def ellipse_in(self, rx: float, ry: float, x0: float, y0: float) -> None: 

218 r''' 

219 Select sites according to 

220 

221 .. math:: 

222 

223 (x-x_0)^2/a^2+(y-y_0)^2/b^2 < 1\, . 

224 

225 :param list_hop: List of Dictionary (see set_hopping definition). 

226 :param rx: Positive Real number. Radius along :math:`x`. 

227 :param ry: Positive Real number. Radius along :math:`y`. 

228 :param x0: Real number. :math:`x` center. 

229 :param y0: Real number. :math:`y` center. 

230 ''' 

231 error_handling.empty_coor(self.coor) 

232 error_handling.positive_real(rx, 'rx') 

233 error_handling.positive_real(ry, 'ry') 

234 error_handling.real_number(x0, 'x0') 

235 error_handling.real_number(y0, 'y0') 

236 self.coor = self.coor[(self.coor['x'] -x0) ** 2 / rx ** 2 + \ 

237 (self.coor['y'] -y0) ** 2 / ry ** 2 < 1.] 

238 self.sites = len(self.coor) 

239 

240 def ellipse_out(self, rx: float, ry: float, x0: float, y0: float) -> None: 

241 r''' 

242 Select sites according to 

243 

244 .. math:: 

245 

246 (x-x_0)^2/a^2+(y-y_0)^2/b^2 > 1\, . 

247 

248 

249 :param list_hop: List of Dictionary (see set_hopping definition). 

250 :param rx: Positive Real number. Radius along :math:`x`. 

251 :param ry: Positive Real number. Radius along :math:`y`. 

252 :param x0: Real number. :math:`x` center. 

253 :param y0: Real number. :math:`y` center. 

254 ''' 

255 error_handling.empty_coor(self.coor) 

256 error_handling.positive_real(rx, 'rx') 

257 error_handling.positive_real(ry, 'ry') 

258 error_handling.real_number(x0, 'x0') 

259 error_handling.real_number(y0, 'y0') 

260 self.coor = self.coor[(self.coor['x'] -x0) ** 2 / rx ** 2 + \ 

261 (self.coor['y'] -y0) ** 2 / ry ** 2 > 1.] 

262 self.sites = len(self.coor) 

263 

264 def center(self) -> None: 

265 ''' 

266 Fix the center of mass of the lattice at (0, 0). 

267 ''' 

268 error_handling.empty_coor(self.coor) 

269 self.coor['x'] -= np.mean(self.coor['x']) 

270 self.coor['y'] -= np.mean(self.coor['y']) 

271 

272 def rotation(self, theta: float) -> None: 

273 r''' 

274 Rotate the lattice structure by the angle :math:`\theta`. 

275 

276 :param theta: Rotation angle in degrees. 

277 ''' 

278 error_handling.empty_coor(self.coor) 

279 error_handling.real_number(theta, 'theta') 

280 theta *= PI / 360 

281 for dic in self.unit_cell: 

282 x = self.coor['x'] - dic['r0'][0] 

283 y = self.coor['y'] - dic['r0'][1] 

284 self.coor['x'] = x * np.cos(theta) - y * np.sin(theta) + dic['r0'][0] 

285 self.coor['y'] = y * np.cos(theta) + x* np.sin(theta) + dic['r0'][1] 

286 

287 def clean_coor(self) -> None: 

288 ''' 

289 Keep only the sites with different coordinates. 

290 ''' 

291 error_handling.empty_coor(self.coor) 

292 coor = self.coor[['x', 'y']].copy() 

293 coor['x'], coor['y'] = self.coor['x'].round(4), self.coor['y'].round(4) 

294 _, idx = np.unique(coor, return_index=True) 

295 self.coor = self.coor[idx] 

296 self.sites = len(self.coor) 

297 

298 def __add__(self, other: 'Lattice') -> 'Lattice': 

299 ''' 

300 Overloading operator +. 

301 ''' 

302 error_handling.lat(other) 

303 error_handling.empty_coor(self.coor) 

304 error_handling.empty_coor(other.coor) 

305 coor = np.concatenate([self.coor, other.coor]) 

306 tags = np.concatenate([self.tags, other.tags]) 

307 lat = lattice(unit_cell=self.unit_cell, prim_vec=self.prim_vec) 

308 lat.add_sites(coor) 

309 lat.sites = self.sites + other.sites 

310 lat.tags = np.unique(tags) 

311 return lat 

312 

313 def __iadd__(self, other: 'Lattice') -> 'Lattice': 

314 ''' 

315 Overloading operator +=. 

316 ''' 

317 error_handling.lat(other) 

318 error_handling.empty_coor(self.coor) 

319 error_handling.empty_coor(other.coor) 

320 self.coor = np.concatenate([self.coor, other.coor]) 

321 self.sites += other.sites 

322 self.tags = np.unique(np.concatenate([self.tags, other.tags])) 

323 return self 

324 

325 def __sub__(self, other: 'Lattice') -> 'Lattice': 

326 ''' 

327 Overloading operator -. 

328 

329 .. note:: 

330 

331 The tags are not considered in the lattice subtraction. 

332 ''' 

333 error_handling.lat(other) 

334 error_handling.empty_coor(self.coor) 

335 error_handling.empty_coor(other.coor) 

336 boo = np.zeros(self.sites, bool) 

337 for c in other.coor: 

338 boo += np.isclose(c['x'], self.coor['x']) & np.isclose(c['y'], self.coor['y']) 

339 coor = self.coor[np.logical_not(boo)] 

340 lat = lattice(unit_cell=self.unit_cell, prim_vec=self.prim_vec) 

341 lat.add_sites(coor) 

342 lat.sites = len(lat.coor) 

343 lat.tags = self.tags 

344 return lat 

345 

346 def __isub__(self, other: 'Lattice') -> 'Lattice': 

347 ''' 

348 Overloading operator -=. 

349 

350 .. note:: 

351 

352 The tags are not considered in the lattice subtraction. 

353 ''' 

354 error_handling.lat(other) 

355 error_handling.empty_coor(self.coor) 

356 error_handling.empty_coor(other.coor) 

357 boo = np.zeros(self.sites, bool) 

358 for c in other.coor: 

359 boo += np.isclose(c['x'], self.coor['x']) & np.isclose(c['y'], self.coor['y']) 

360 self.coor = self.coor[np.logical_not(boo)] 

361 self.sites = int(np.sum(np.logical_not(boo))) 

362 return self 

363 

364 def plot( 

365 self, 

366 ms: float = 20, 

367 fs: float = 20, 

368 plt_index: bool = False, 

369 axis: bool = False, 

370 figsize: tuple[float, float] | None = None, 

371 ) -> Figure: 

372 ''' 

373 Plot lattice in hopping space. 

374 

375 :param ms: Positive number. Default value 20. Markersize. 

376 :param fs: Positve number. Default value 20. Fontsize. 

377 :param plt_index: Boolean. Default value False. Plot site labels. 

378 :param axis: Boolean. Default value False. Plot axis. 

379 :param figsize: Tuple. Default value None. Figsize. 

380 

381 :returns: 

382 * **fig** -- Figure. 

383 ''' 

384 error_handling.empty_coor(self.coor) 

385 error_handling.positive_real(ms, 'ms') 

386 error_handling.positive_real(fs, 'fs') 

387 error_handling.boolean(plt_index, 'plt_index') 

388 error_handling.boolean(axis, 'axis') 

389 if figsize is None: 

390 figsize = (5, 5) 

391 error_handling.list_tuple_2elem(figsize, 'figsize') 

392 error_handling.positive_real(figsize[0], 'figsize[0]') 

393 error_handling.positive_real(figsize[1], 'figsize[1]') 

394 fig, ax = plt.subplots(figsize=figsize) 

395 # plot sites 

396 colors = ['b', 'r', 'g', 'y', 'm', 'k'] 

397 for color, tag in zip(colors, self.tags): 

398 plt.plot(self.coor['x'][self.coor['tag'] == tag], 

399 self.coor['y'][self.coor['tag'] == tag], 

400 'o', color=color, ms=ms, markeredgecolor='none') 

401 ax.set_aspect('equal') 

402 ax.set_xlim([np.min(self.coor['x'])-1., np.max(self.coor['x'])+1.]) 

403 ax.set_ylim([np.min(self.coor['y'])-1., np.max(self.coor['y'])+1.]) 

404 if not axis: 

405 ax.axis('off') 

406 # plot indices 

407 if plt_index: 

408 indices = ['{}'.format(i) for i in range(self.sites)] 

409 for l, x, y in zip(indices, self.coor['x'], self.coor['y']): 

410 plt.annotate(l, xy=(x, y), xytext=(0, 0), 

411 textcoords='offset points', 

412 ha='right', va='bottom', size=fs) 

413 plt.draw() 

414 return fig 

415 

416 def show(self) -> None: 

417 """ 

418 Emulate Matplotlib method plt.show(). 

419 """ 

420 plt.show() 

421 

422 

423# Backward-compatible lowercase alias (pre-0.2 API). 

424lattice = Lattice