Coverage for tbkit/plot.py: 100%

1from __future__ import annotations 

2 

3import numpy as np 

4from numpy.typing import NDArray, ArrayLike 

5import matplotlib.pyplot as plt 

6from matplotlib.figure import Figure 

7from matplotlib.axes import Axes 

8from mpl_toolkits.mplot3d import Axes3D 

9from matplotlib.legend_handler import HandlerLine2D 

10import tbkit.error_handling as error_handling 

11import tbkit.dos as dos 

12from tbkit.system import System 

13import os 

14 

15 

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

17# CLASS PLOT 

18################################# 

19 

20 

21class Plot: 

22 ''' 

23 Plot the results of the classes **lattice** or **system**. 

24 

25 :param sys: class instance **system**. 

26 :param colors: Default value None. Color plot. 

27 ''' 

28 

29 def __init__(self, sys: System, colors: list[str] | None = None) -> None: 

30 error_handling.sys(sys) 

31 self.sys = sys 

32 if colors is None: 

33 self.colors = ['b', 'r', 'g', 'y', 'm', 'k'] 

34 else: 

35 self.colors = colors 

36 

37 def plt_hopping(self, coor: NDArray, hop: NDArray, c: float) -> None: 

38 ''' 

39 Private method called by *lattice_generic*. 

40 ''' 

41 for i in range(len(hop)): 

42 plt.plot([coor['x'][hop['i'][i]], 

43 coor['x'][hop['j'][i]]], 

44 [coor['y'][hop['i'][i]], 

45 coor['y'][hop['j'][i]]], 

46 'k', lw=c*hop['t'][i].real) 

47 

48 def lattice_generic( 

49 self, 

50 coor: NDArray, 

51 ms: float, 

52 lw: float, 

53 c: float, 

54 fs: float, 

55 axis: bool, 

56 plt_hop: bool, 

57 plt_hop_low: bool, 

58 plt_index: bool, 

59 figsize: tuple[float, float] | None, 

60 ) -> Figure: 

61 ''' 

62 Private method called by *lattice* and *lattice_hop*. 

63 ''' 

64 error_handling.positive_real(ms, 'ms') 

65 error_handling.positive_real(lw, 'lw') 

66 error_handling.positive_real(c, 'c') 

67 error_handling.positive_real(fs, 'fs') 

68 error_handling.boolean(axis, 'axis') 

69 error_handling.boolean(plt_hop, 'plt_hop') 

70 error_handling.boolean(plt_hop_low, 'plt_hop_low') 

71 error_handling.boolean(plt_index, 'plt_index') 

72 error_handling.tuple_2elem(figsize, 'figsize') 

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

74 # hoppings 

75 if plt_hop: 

76 error_handling.empty_ndarray(self.sys.hop, 'sys.hop') 

77 self.plt_hopping(coor, self.sys.hop[self.sys.hop['ang']>=0], c) 

78 if plt_hop_low: 

79 error_handling.empty_ndarray(self.sys.hop, 'sys.hop') 

80 self.plt_hopping(coor, self.sys.hop[self.sys.hop['ang']<0], c) 

81 # plot sites 

82 for color, tag in zip(self.colors, self.sys.lat.tags): 

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

84 coor['y'][coor['tag'] == tag], 

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

86 ax.set_aspect('equal') 

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

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

89 if not axis: 

90 ax.axis('off') 

91 # plot indices 

92 if plt_index: 

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

94 for l, x, y in zip(indices, coor['x'], coor['y']): 

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

96 textcoords='offset points', 

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

98 plt.draw() 

99 return fig 

100 

101 def lattice( 

102 self, 

103 ms: float = 20, 

104 lw: float = 5., 

105 c: float = 3., 

106 fs: float = 20, 

107 axis: bool = False, 

108 plt_hop: bool = False, 

109 plt_hop_low: bool = False, 

110 plt_index: bool = False, 

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

112 ) -> Figure: 

113 ''' 

114 Plot lattice. 

115 

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

117 :param c: Positive number. Default value 3.  

118 Coefficient. Hopping linewidths given by c*hop['t']. 

119 :param fs: Positive number. Default value 20. Fontsize. 

120 :param plt_hop: Boolean. Default value False. Plot hoppings. 

121 :param plt_hop_low: Boolean. Default value False.  

122 Plot hoppings diagonal low. 

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

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

125 :param figsize: Tuple. Default value None. Figure size. 

126 

127 :returns: 

128 * **fig** -- Figure. 

129 ''' 

130 error_handling.empty_ndarray(self.sys.lat.coor, 'lat.get_lattice') 

131 return self.lattice_generic(self.sys.lat.coor, ms, lw, c, fs, axis, plt_hop, 

132 plt_hop_low, plt_index, figsize) 

133 

134 def lattice_hop( 

135 self, 

136 ms: float = 20, 

137 lw: float = 5, 

138 c: float = 3., 

139 fs: float = 20, 

140 axis: bool = False, 

141 plt_hop: bool = False, 

142 plt_hop_low: bool = False, 

143 plt_index: bool = False, 

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

145 ) -> Figure: 

146 ''' 

147 Plot lattice in hopping space. 

148 

149 :param ms: Positive Float. Default value 20. Markersize. 

150 :param c: Positive Float. Default value 3. Coefficient. 

151 Hopping linewidths given by c*hop['t']. 

152 :param fs: Positive Float. Default value 20. Fontsize. 

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

154 :param plt_hop: Boolean. Default value False. Plot hoppings. 

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

156 :param figsize: Tuple. Default value None. Figure size. 

157 

158 :returns: 

159 * **fig** -- Figure. 

160 ''' 

161 error_handling.empty_ndarray(self.sys.coor_hop, 'sys.get_coor_hop') 

162 return self.lattice_generic(self.sys.coor_hop, ms, lw, c, fs, axis, plt_hop, 

163 plt_hop_low, plt_index, figsize) 

164 

165 

166 def spectrum_hist(self, nbr_bins: int = 61, fs: float = 20, lims: tuple[float, float] | None = None) -> None: 

167 """ 

168 Plot the spectrum. 

169  

170 :param nbr_bins: Default value 101. Number of bins of the histogram. 

171 :param lims: List, lims[0] energy min, lims[1] energy max. 

172 """ 

173 error_handling.empty_ndarray(self.sys.en, 'sys.get_eig') 

174 error_handling.positive_real(nbr_bins, 'nbr_bins') 

175 error_handling.lims(lims) 

176 fig, ax = plt.subplots() 

177 if lims is None: 

178 en_max = np.max(self.sys.en.real) 

179 ind_en = np.ones(self.sys.lat.sites, bool) 

180 ax.set_ylim([-en_max, en_max]) 

181 else: 

182 ind_en = np.argwhere((self.sys.en > lims[0]) & (self.sys.en < lims[1])) 

183 ind_en = np.ravel(ind_en) 

184 ax.set_xlim(lims) 

185 en = self.sys.en[ind_en] 

186 n, bins, patches = plt.hist(en, bins=nbr_bins, color='b', alpha=0.8) 

187 ax.set_title('Spectrum', fontsize=fs) 

188 ax.set_xlabel('$E$', fontsize=fs) 

189 ax.set_ylabel('number of states', fontsize=fs) 

190 ax.set_ylim([0, np.max(n)+1]) 

191 for label in ax.xaxis.get_majorticklabels(): 

192 label.set_fontsize(fs) 

193 for label in ax.yaxis.get_majorticklabels(): 

194 label.set_fontsize(fs) 

195 

196 def dos( 

197 self, 

198 broadening: float = 0.05, 

199 kernel: str = 'gaussian', 

200 e_grid: ArrayLike | None = None, 

201 fs: float = 20, 

202 lw: float = 2., 

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

204 ) -> Figure: 

205 ''' 

206 Plot the (broadened) density of states, see *tbkit.dos.density_of_states*. 

207 

208 :param broadening: Positive real number. Default value 0.05. Kernel width. 

209 :param kernel: String. Default value 'gaussian'. 'gaussian' or 'lorentzian'. 

210 :param e_grid: Real ndarray. Default value None. Energies at which to 

211 evaluate the density of states. 

212 :param fs: Positive number. Default value 20. Fontsize. 

213 :param lw: Positive number. Default value 2. Linewidth. 

214 :param figsize: Tuple. Default value None. Figure size. 

215 

216 :returns: 

217 * **fig** -- Figure. 

218 ''' 

219 error_handling.empty_ndarray(self.sys.en, 'sys.get_eig') 

220 error_handling.positive_real(fs, 'fs') 

221 error_handling.positive_real(lw, 'lw') 

222 error_handling.tuple_2elem(figsize, 'figsize') 

223 e_grid, rho = dos.density_of_states(self.sys.en, e_grid=e_grid, 

224 broadening=broadening, kernel=kernel) 

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

226 ax.plot(e_grid, rho, 'b', lw=lw) 

227 ax.fill_between(e_grid, rho, color='b', alpha=0.2) 

228 ax.set_xlim([e_grid[0], e_grid[-1]]) 

229 ax.set_ylim([0., None]) 

230 ax.set_title('Density of states', fontsize=fs) 

231 ax.set_xlabel('$E$', fontsize=fs) 

232 ax.set_ylabel(r'$\rho(E)$', fontsize=fs) 

233 for label in ax.xaxis.get_majorticklabels(): 

234 label.set_fontsize(fs) 

235 for label in ax.yaxis.get_majorticklabels(): 

236 label.set_fontsize(fs) 

237 fig.set_layout_engine('tight') 

238 plt.draw() 

239 return fig 

240 

241 def spectrum( 

242 self, 

243 ms: float = 10, 

244 fs: float = 20, 

245 lims: tuple[float, float] | None = None, 

246 tag_pola: str | None = None, 

247 ipr: bool | None = None, 

248 peterman: bool | None = None, 

249 ) -> Figure: 

250 ''' 

251 Plot spectrum (eigenenergies real part (blue circles), 

252 and sublattice polarization if *pola* not empty (red circles). 

253 

254 :param ms: Default value 10. Markersize. 

255 :param fs: Default value 20. Fontsize. 

256 :param lims: List, lims[0] energy min, lims[1] energy max. 

257 :param tag_pola: Default value None. One-character string. Tag of the sublattice. 

258 :param ipr: Default value None. If True plot the Inverse Partitipation Ration. 

259 :param petermann: Default value None. If True plot the Petermann factor. 

260 

261 :returns: 

262 * **fig** -- Figure. 

263 ''' 

264 error_handling.empty_ndarray(self.sys.en, 'sys.get_eig') 

265 error_handling.positive_real(ms, 'ms') 

266 error_handling.positive_real(fs, 'fs') 

267 error_handling.lims(lims) 

268 fig, ax1 = plt.subplots() 

269 ax1 = plt.gca() 

270 x = np.arange(self.sys.lat.sites) 

271 if lims is None: 

272 en_max = np.max(self.sys.en.real) 

273 ax1.set_ylim([-en_max-0.2, en_max+0.2]) 

274 ind = np.ones(self.sys.lat.sites, bool) 

275 else: 

276 ind = (self.sys.en > lims[0]) & (self.sys.en < lims[1]) 

277 ax1.set_ylim([lims[0]-0.1, lims[1]+0.1]) 

278 ax1.plot(x[ind], self.sys.en.real[ind], 'ob', markersize=ms) 

279 ax1.set_title('Spectrum', fontsize=fs) 

280 ax1.set_xlabel('$n$', fontsize=fs) 

281 ax1.set_ylabel('$E_n$', fontsize=fs, color='blue') 

282 for label in ax1.get_yticklabels(): 

283 label.set_color('b') 

284 if tag_pola: 

285 error_handling.tag(tag_pola, self.sys.lat.tags) 

286 fig, ax2 = self.polarization(fig=fig, ax1=ax1, ms=ms, fs=fs, tag_pola=tag_pola, ind=ind) 

287 elif ipr: 

288 fig, ax2 = self.ipr(fig=fig, ax1=ax1, ms=ms, fs=fs, ind=ind) 

289 elif peterman: 

290 fig, ax2 = self.petermann(fig=fig, ax1=ax1, ms=ms, fs=fs, ind=ind) 

291 for label in ax1.xaxis.get_majorticklabels(): 

292 label.set_fontsize(fs) 

293 for label in ax1.yaxis.get_majorticklabels(): 

294 label.set_fontsize(fs) 

295 xa = ax1.get_xaxis() 

296 ax1.set_xlim([x[ind][0]-0.5, x[ind][-1]+0.5]) 

297 xa.set_major_locator(plt.MaxNLocator(integer=True)) 

298 fig.set_layout_engine('tight') 

299 plt.draw() 

300 return fig 

301 

302 def polarization( 

303 self, 

304 fig: Figure | None = None, 

305 ax1: Axes | None = None, 

306 ms: float = 10., 

307 fs: float = 20., 

308 lims: tuple[float, float] | None = None, 

309 tag_pola: str | None = None, 

310 ind: NDArray | None = None, 

311 ) -> tuple[Figure, Axes]: 

312 ''' 

313 Plot sublattice polarization. 

314 

315 :param fig: Figure. Default value None. (used by the method spectrum). 

316 :param ax1: Axis. Default value None. (used by the method spectrum). 

317 :param ms: Positive Float. Default value 10. Markersize. 

318 :param fs: Positive Float. Default value 20. Fontsize. 

319 :param lims: List, lims[0] energy min, lims[1] energy max. 

320 :param tag_pola: One-character string. Default value None. Tag of the sublattice. 

321 :param ind: List. Default value None. List of indices. (used in the method spectrum). 

322 

323 :returns: 

324 * **fig** -- Figure. 

325 ''' 

326 if fig is None: 

327 error_handling.sys(self.sys) 

328 error_handling.empty_ndarray(self.sys.en, 'sys.get_eig') 

329 error_handling.positive_real(ms, 'ms') 

330 error_handling.positive_real(fs, 'fs') 

331 error_handling.lims(lims) 

332 fig, ax2 = plt.subplots() 

333 ax2 = plt.gca() 

334 if lims is None: 

335 ax2.set_ylim([-0.1, 1.1]) 

336 ind = np.ones(self.sys.lat.sites, bool) 

337 else: 

338 ind = (self.sys.en > lims[0]) & (self.sys.en < lims[1]) 

339 ax2.set_ylim([lims[0]-0.1, lims[1]+0.1]) 

340 else: 

341 ax2 = plt.twinx() 

342 error_handling.empty_ndarray(self.sys.pola, 'sys.get_pola') 

343 error_handling.tag(tag_pola, self.sys.lat.tags) 

344 x = np.arange(self.sys.lat.sites) 

345 i_tag = self.sys.lat.tags == tag_pola 

346 ax2.plot(x[ind], np.ravel(self.sys.pola[ind, i_tag]), 'or', markersize=(4*ms)//5) 

347 ylabel = '$<' + tag_pola.upper() + '|' + tag_pola.upper() + '>$' 

348 ax2.set_ylabel(ylabel, fontsize=fs, color='red') 

349 ax2.set_ylim([-0.1, 1.1]) 

350 ax2.set_xlim(-0.5, x[ind][-1]+0.5) 

351 for tick in ax2.xaxis.get_major_ticks(): 

352 tick.label1.set_fontsize(fs) 

353 for label in ax2.get_yticklabels(): 

354 label.set_color('r') 

355 return fig, ax2 

356 

357 def ipr( 

358 self, 

359 fig: Figure | None = None, 

360 ax1: Axes | None = None, 

361 ms: float = 10, 

362 fs: float = 20, 

363 lims: tuple[float, float] | None = None, 

364 ind: NDArray | None = None, 

365 ) -> tuple[Figure, Axes]: 

366 ''' 

367 Plot Inverse Participation Ration. 

368 

369 :param fig: Figure. Default value None. (used by the method spectrum). 

370 :param ax1: Axis. Default value None. (used by the method spectrum). 

371 :param ms: Positive Float. Default value 10. Markersize. 

372 :param fs: Positive Float. Default value 20. Fontsize. 

373 :param lims: List. lims[0] energy min, lims[1] energy max. 

374 :param ind: List. Default value None. List of indices. (used in the method spectrum). 

375 

376 :returns: 

377 * **fig** -- Figure. 

378 ''' 

379 if fig is None: 

380 error_handling.sys(self.sys) 

381 error_handling.empty_ndarray(self.sys.ipr, 'sys.get_ipr') 

382 error_handling.positive_real(ms, 'ms') 

383 error_handling.positive_real(fs, 'fs') 

384 error_handling.lims(lims) 

385 fig, ax2 = plt.subplots() 

386 ax2 = plt.gca() 

387 if lims is None: 

388 ind = np.ones(self.sys.lat.sites, bool) 

389 else: 

390 ind = (self.sys.en > lims[0]) & (self.sys.en < lims[1]) 

391 else: 

392 ax2 = plt.twinx() 

393 error_handling.empty_ndarray(self.sys.ipr, 'sys.get_ipr') 

394 x = np.arange(self.sys.lat.sites) 

395 ax2.plot(x[ind], self.sys.ipr[ind], 'or', markersize=(4*ms)//5) 

396 ax2.set_ylabel( 'IPR' , fontsize=fs, color='red') 

397 ax2.set_xlim(-0.5, x[ind][-1]+0.5) 

398 for tick in ax2.xaxis.get_major_ticks(): 

399 tick.label1.set_fontsize(fs) 

400 for label in ax2.get_yticklabels(): 

401 label.set_color('r') 

402 return fig, ax2 

403 

404 def petermann( 

405 self, 

406 fig: Figure | None = None, 

407 ax1: Axes | None = None, 

408 ms: float = 10, 

409 fs: float = 20, 

410 lims: tuple[float, float] | None = None, 

411 ind: NDArray | None = None, 

412 ) -> tuple[Figure, Axes]: 

413 ''' 

414 Plot Peterman factor. 

415 

416 :param fig: Figure. Default value None. (used by the method spectrum). 

417 :param ax1: Axis. Default value None. (used by the method spectrum). 

418 :param ms: Positive Float. Default value 10. Markersize. 

419 :param fs: Positive Float. Default value 20. Fontsize. 

420 :param lims: List. lims[0] energy min, lims[1] energy max. 

421 :param ind: List. Default value None. List of indices. (used in the method spectrum). 

422 

423 :returns: 

424 * **fig** -- Figure. 

425 ''' 

426 if fig is None: 

427 error_handling.sys(self.sys) 

428 error_handling.empty_ndarray(self.sys.ipr, 'sys.get_petermann') 

429 error_handling.positive_real(ms, 'ms') 

430 error_handling.positive_real(fs, 'fs') 

431 error_handling.lims(lims) 

432 fig, ax2 = plt.subplots() 

433 ax2 = plt.gca() 

434 if lims is None: 

435 ind = np.ones(self.sys.lat.sites, bool) 

436 else: 

437 ind = (self.sys.en > lims[0]) & (self.sys.en < lims[1]) 

438 else: 

439 ax2 = plt.twinx() 

440 error_handling.empty_ndarray(self.sys.ipr, 'sys.get_ipr') 

441 x = np.arange(self.sys.lat.sites) 

442 ax2.plot(x[ind], self.sys.petermann[ind], 'or', markersize=(4*ms)//5) 

443 ax2.set_ylabel( 'K' , fontsize=fs, color='red') 

444 ax2.set_xlim(-0.5, x[ind][-1]+0.5) 

445 for tick in ax2.xaxis.get_major_ticks(): 

446 tick.label1.set_fontsize(fs) 

447 for label in ax2.get_yticklabels(): 

448 label.set_color('r') 

449 return fig, ax2 

450 

451 def spectrum_complex(self, ms: float = 10., fs: float = 20., lims: tuple[float, float] | None = None) -> Figure: 

452 ''' 

453 Plot complex value eigenenergies, real part (blue circles), 

454 and imaginary part (red circles). 

455 

456 :param ms: Positive Float. Default value 20. Markersize. 

457 :param fs: Positive Float. Default value 20. Font size. 

458 :param lims: List. lims[0] energy min, lims[1] energy max. 

459 

460 :returns: 

461 * **fig** -- Figure. 

462 ''' 

463 error_handling.empty_ndarray(self.sys.en, 'sys.get_eig') 

464 error_handling.positive_real(ms, 'ms') 

465 error_handling.positive_real(fs, 'fs') 

466 error_handling.lims(lims) 

467 fig, ax1 = plt.subplots() 

468 ax1 = plt.gca() 

469 x = np.arange(self.sys.lat.sites) 

470 if lims is None: 

471 en_max = np.max(self.sys.en.real) 

472 ax1.set_ylim([-en_max-0.2, en_max+0.2]) 

473 ind = np.ones(self.sys.lat.sites, bool) 

474 else: 

475 ind = (self.sys.en > lims[0]) & (self.sys.en < lims[1]) 

476 ax1.set_ylim([lims[0]-0.1, lims[1]+0.1]) 

477 ax1.plot(x[ind], self.sys.en.real[ind], 'ob', markersize=ms) 

478 ax1.plot(x[ind], self.sys.en.imag[ind], 'or', markersize=ms) 

479 ax1.set_title('Spectrum', fontsize=fs) 

480 ax1.set_xlabel('$n$', fontsize=fs) 

481 ax1.set_ylabel('Re '+r'$E_n$'+', Im '+r'$E_n$', fontsize=fs) 

482 for label in ax1.xaxis.get_majorticklabels(): 

483 label.set_fontsize(fs) 

484 for label in ax1.yaxis.get_majorticklabels(): 

485 label.set_fontsize(fs) 

486 xa = ax1.get_xaxis() 

487 ax1.set_xlim([x[ind][0]-0.1, x[ind][-1]+0.1]) 

488 xa.set_major_locator(plt.MaxNLocator(integer=True)) 

489 fig.set_layout_engine('tight') 

490 plt.draw() 

491 return fig 

492 

493 def intensity_1d( 

494 self, intensity: NDArray, ms: float = 20., lw: float = 2., fs: float = 20., 

495 title: str = r'$|\psi^{(j)}|^2$', 

496 ) -> Figure: 

497 ''' 

498 Plot intensity for 1D lattices. 

499 

500 :param intensity: np.array. Field intensity. 

501 :param ms: Positive Float. Default value 20. Markersize. 

502 :param lw: Positive Float. Default value 2. Linewith, connect sublattice sites. 

503 :param fs: Positive Float. Default value 20. Font size. 

504 :param title: String. Default value 'Intensity'. Figure title. 

505 ''' 

506 error_handling.ndarray(intensity, 'intensity', self.sys.lat.sites) 

507 error_handling.empty_ndarray(self.sys.lat.coor, 'sys.get_lattice') 

508 error_handling.positive_real(ms, 'ms') 

509 error_handling.positive_real(lw, 'lw') 

510 error_handling.positive_real(fs, 'fs') 

511 error_handling.string(title, 'title') 

512 fig, ax = plt.subplots() 

513 ax.set_xlabel('$j$', fontsize=fs) 

514 ax.set_ylabel(title, fontsize=fs) 

515 ax.set_title(title, fontsize=fs) 

516 for t, c in zip(self.sys.lat.tags, self.colors): 

517 plt.plot(self.sys.lat.coor['x'][self.sys.lat.coor['tag'] == t], 

518 intensity[self.sys.lat.coor['tag'] == t], 

519 '-o', color=c, ms=ms, lw=lw) 

520 plt.xlim([-1., self.sys.lat.sites]) 

521 plt.ylim([0., np.max(intensity)+.05]) 

522 fig.set_layout_engine('tight') 

523 plt.draw() 

524 return fig 

525 

526 def intensity_disk( 

527 self, 

528 intensity: NDArray, 

529 s: float = 200., 

530 fs: float = 20., 

531 lims: tuple[float, float] | None = None, 

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

533 title: str = r'$|\psi|^2$', 

534 ) -> Figure: 

535 r''' 

536 Plot the intensity. Colormap with identical disk shape. 

537 

538 :param intensity: np.array.Field intensity. 

539 :param s: Default value 200. Disk size. 

540 :param fs: Default value 20. Font size. 

541 :param lims: List. Default value None. Colormap limits. 

542 :param figsize: Tuple. Default value None. Figure size. 

543 :param title: String. Default value '$|\psi_n|^2$'. Title. 

544 

545 :returns: 

546 * **fig** -- Figure. 

547 ''' 

548 error_handling.empty_ndarray(self.sys.lat.coor, 'sys.get_lattice') 

549 error_handling.ndarray(intensity, 'intensity', self.sys.lat.sites) 

550 error_handling.positive_real(s, 's') 

551 error_handling.positive_real(fs, 'fs') 

552 error_handling.tuple_2elem(figsize, 'figsize') 

553 error_handling.string(title, 'title') 

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

555 plt.title(title, fontsize=fs+5) 

556 map_red = plt.get_cmap('Reds') 

557 if lims is None: 

558 lims = [0., np.max(intensity)] 

559 y_ticks = ['0', 'max'] 

560 else: 

561 y_ticks = lims 

562 plt.scatter(self.sys.lat.coor['x'], self.sys.lat.coor['y'], c=intensity, s=s, 

563 cmap=map_red, vmin=lims[0], vmax=lims[1]) 

564 cbar = plt.colorbar(ticks=lims) 

565 ax.set_xticks([]) 

566 ax.set_yticks([]) 

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

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

569 cbar.ax.set_yticklabels([y_ticks[0], y_ticks[1]]) 

570 cbar.ax.tick_params(labelsize=fs) 

571 ax.set_aspect('equal') 

572 fig.set_layout_engine('tight') 

573 plt.draw() 

574 return fig 

575 

576 def intensity_area( 

577 self, 

578 intensity: NDArray, 

579 s: float = 1000., 

580 lw: float = 1., 

581 fs: float = 20., 

582 plt_hop: bool = False, 

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

584 title: str = r'$|\psi|^2$', 

585 ) -> Figure: 

586 r''' 

587 Plot the intensity. Intensity propotional to disk shape. 

588 

589 :param intensity: np.array. Intensity. 

590 :param s: Positive Float. Default value 1000. 

591 Circle size given by s * intensity. 

592 :param lw: Positive Float. Default value 1. Hopping linewidths. 

593 :param fs: Positive Float. Default value 20. Fontsize. 

594 :param plt_hop: Boolean. Default value False. Plot hoppings. 

595 :param figsize: Tuple. Default value None. Figure size. 

596 :param title: String. Default value '$|\psi_{ij}|^2$'. Figure title. 

597 

598 :returns: 

599 * **fig** -- Figure. 

600 ''' 

601 error_handling.empty_ndarray(self.sys.lat.coor, 'sys.get_lattice') 

602 error_handling.ndarray(intensity, 'intensity', self.sys.lat.sites) 

603 error_handling.positive_real(s, 's') 

604 error_handling.positive_real(fs, 'fs') 

605 error_handling.boolean(plt_hop, 'plt_hop') 

606 error_handling.tuple_2elem(figsize, 'figsize') 

607 error_handling.string(title, 'title') 

608 fig, ax = plt.subplots() 

609 ax.set_xlabel('$i$', fontsize=fs) 

610 ax.set_ylabel('$j$', fontsize=fs) 

611 ax.set_title(title, fontsize=fs) 

612 if plt_hop: 

613 plt.plot([self.sys.lat.coor['x'][self.sys.hop['i'][:]], 

614 self.sys.lat.coor['x'][self.sys.hop['j'][:]]], 

615 [self.sys.lat.coor['y'][self.sys.hop['i'][:]], 

616 self.sys.lat.coor['y'][self.sys.hop['j'][:]]], 

617 'k', lw=lw) 

618 for tag, color in zip(self.sys.lat.tags, self.colors): 

619 plt.scatter(self.sys.lat.coor['x'][self.sys.lat.coor['tag'] == tag], 

620 self.sys.lat.coor['y'][self.sys.lat.coor['tag'] == tag], 

621 s=100*s*intensity[self.sys.lat.coor['tag'] == tag], 

622 c=color, alpha=0.5) 

623 ax.set_aspect('equal') 

624 ax.axis('off') 

625 x_lim = [np.min(self.sys.lat.coor['x'])-2., np.max(self.sys.lat.coor['x'])+2.] 

626 y_lim = [np.min(self.sys.lat.coor['y'])-2., np.max(self.sys.lat.coor['y'])+2.] 

627 ax.set_xlim(x_lim) 

628 ax.set_ylim(y_lim) 

629 fig.set_layout_engine('tight') 

630 plt.draw() 

631 return fig 

632 

633 def butterfly( 

634 self, 

635 betas: NDArray, 

636 butterfly: NDArray, 

637 lw: float = 1., 

638 fs: float = 20., 

639 lims: tuple[float, float] | None = None, 

640 title: str = '', 

641 ) -> Figure: 

642 ''' 

643 Plot energies depending on a parameter. 

644 

645 :param betas: np.array. Parameter values. 

646 :param butterfly: np.array. Eigenvalues. 

647 :param lw: Positive Float. Default value 1. Hopping linewidths. 

648 :param fs: Positive Float. Default value 20. Fontsize. 

649 :param lims: List, lims[0] energy min, lims[1] energy max. 

650 :param title: Default value ''. Figure title. 

651 ''' 

652 error_handling.ndarray_empty(betas, 'betas') 

653 error_handling.ndarray_empty(butterfly, 'butterfly')