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, SubFigure 

7from matplotlib.axes import Axes 

8import tbkit.error_handling as error_handling 

9import tbkit.dos as dos 

10from tbkit.system import System 

11 

12 

13################################# 

14# CLASS PLOT 

15################################# 

16 

17 

18class Plot: 

19 ''' 

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

21 

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

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

24 ''' 

25 

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

27 error_handling.sys(sys) 

28 self.sys = sys 

29 if colors is None: 

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

31 else: 

32 self.colors = colors 

33 

34 def plt_hopping(self, coor: NDArray, hop: NDArray, c: float, ax: Axes | None = None) -> None: 

35 ''' 

36 Private method called by *lattice_generic*. 

37 ''' 

38 if ax is None: 

39 ax = plt.gca() 

40 # bonds that wrap around a torus would cross the whole sample 

41 hop = hop[~self.sys._wrapped(hop['i'].astype(int), hop['j'].astype(int))] 

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

43 ax.plot([coor['x'][hop['i'][i]], 

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

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

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

47 'k', lw=c*np.abs(hop['t'][i])) 

48 

49 def lattice_generic( 

50 self, 

51 coor: NDArray, 

52 ms: float, 

53 lw: float, 

54 c: float, 

55 fs: float, 

56 axis: bool, 

57 plt_hop: bool, 

58 plt_hop_low: bool, 

59 plt_index: bool, 

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

61 ax: Axes | None = None, 

62 ) -> Figure | SubFigure: 

63 ''' 

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

65 ''' 

66 error_handling.positive_real(ms, 'ms') 

67 error_handling.positive_real(lw, 'lw') 

68 error_handling.positive_real(c, 'c') 

69 error_handling.positive_real(fs, 'fs') 

70 error_handling.boolean(axis, 'axis') 

71 error_handling.boolean(plt_hop, 'plt_hop') 

72 error_handling.boolean(plt_hop_low, 'plt_hop_low') 

73 error_handling.boolean(plt_index, 'plt_index') 

74 error_handling.tuple_2elem(figsize, 'figsize') 

75 fig: Figure | SubFigure 

76 if ax is None: 

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

78 else: 

79 fig = ax.figure 

80 # hoppings 

81 if plt_hop: 

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

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

84 if plt_hop_low: 

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

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

87 # plot sites 

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

89 ax.plot(coor['x'][coor['tag'] == tag], 

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

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

92 ax.set_aspect('equal') 

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

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

95 if not axis: 

96 ax.axis('off') 

97 # plot indices 

98 if plt_index: 

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

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

101 ax.annotate(l, xy=(x, y), xytext=(0, 0), 

102 textcoords='offset points', 

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

104 plt.draw() 

105 return fig 

106 

107 def lattice( 

108 self, 

109 ms: float = 20, 

110 lw: float = 5., 

111 c: float = 3., 

112 fs: float = 20, 

113 axis: bool = False, 

114 plt_hop: bool = False, 

115 plt_hop_low: bool = False, 

116 plt_index: bool = False, 

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

118 ax: Axes | None = None, 

119 ) -> Figure | SubFigure: 

120 ''' 

121 Plot lattice. 

122 

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

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

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

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

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

128 :param plt_hop_low: Boolean. Default value False.  

129 Plot hoppings diagonal low. 

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

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

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

133 :param ax: Axes. Default value None. Draw on this axis instead of 

134 creating a figure of its own -- to place several lattices side 

135 by side in one figure. *figsize* is then ignored. 

136 

137 :returns: 

138 * **fig** -- Figure. 

139 ''' 

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

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

142 plt_hop_low, plt_index, figsize, ax) 

143 

144 def lattice_hop( 

145 self, 

146 ms: float = 20, 

147 lw: float = 5, 

148 c: float = 3., 

149 fs: float = 20, 

150 axis: bool = False, 

151 plt_hop: bool = False, 

152 plt_hop_low: bool = False, 

153 plt_index: bool = False, 

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

155 ax: Axes | None = None, 

156 ) -> Figure | SubFigure: 

157 ''' 

158 Plot lattice in hopping space. 

159 

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

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

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

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

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

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

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

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

168 :param ax: Axes. Default value None. Draw on this axis instead of 

169 creating a figure of its own -- to place several lattices side 

170 by side in one figure. *figsize* is then ignored. 

171 

172 :returns: 

173 * **fig** -- Figure. 

174 ''' 

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

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

177 plt_hop_low, plt_index, figsize, ax) 

178 

179 

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

181 """ 

182 Plot the spectrum (real part of the eigenenergies) as a histogram. 

183 

184 :param nbr_bins: Positive integer. Default value 61. Number of bins of the histogram. 

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

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

187 

188 :returns: 

189 * **fig** -- Figure. 

190 """ 

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

192 error_handling.positive_int(nbr_bins, 'nbr_bins') 

193 error_handling.positive_real(fs, 'fs') 

194 error_handling.lims(lims) 

195 fig, ax = plt.subplots() 

196 en_real = self.sys.en.real 

197 if lims is None: 

198 ind_en = np.ones(len(self.sys.en), bool) 

199 else: 

200 ind_en = np.ravel(np.argwhere((en_real > lims[0]) & (en_real < lims[1]))) 

201 ax.set_xlim(lims) 

202 en = en_real[ind_en] 

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

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

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

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

207 ax.set_ylim(0, np.max(n)+1) 

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

209 label.set_fontsize(fs) 

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

211 label.set_fontsize(fs) 

212 return fig 

213 

214 def dos( 

215 self, 

216 broadening: float = 0.05, 

217 kernel: str = 'gaussian', 

218 e_grid: ArrayLike | None = None, 

219 fs: float = 20, 

220 lw: float = 2., 

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

222 ) -> Figure: 

223 ''' 

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

225 

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

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

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

229 evaluate the density of states. 

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

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

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

233 

234 :returns: 

235 * **fig** -- Figure. 

236 ''' 

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

238 error_handling.positive_real(fs, 'fs') 

239 error_handling.positive_real(lw, 'lw') 

240 error_handling.tuple_2elem(figsize, 'figsize') 

241 return dos._plot_density_of_states(self.sys.en, e_grid, broadening, kernel, fs, lw, figsize) 

242 

243 def spectrum( 

244 self, 

245 ms: float = 10, 

246 fs: float = 20, 

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

248 tag_pola: str | None = None, 

249 ipr: bool | None = None, 

250 petermann: bool | None = None, 

251 ) -> Figure: 

252 ''' 

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

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

255 

256 :param ms: Default value 10. Markersize. 

257 :param fs: Default value 20. Fontsize. 

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

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

260 :param ipr: Default value None. If True plot the Inverse Participation Ratio. 

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

262 

263 :returns: 

264 * **fig** -- Figure. 

265 ''' 

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

267 error_handling.positive_real(ms, 'ms') 

268 error_handling.positive_real(fs, 'fs') 

269 error_handling.lims(lims) 

270 fig, ax1 = plt.subplots() 

271 ax1 = plt.gca() 

272 x = np.arange(len(self.sys.en)) 

273 if lims is None: 

274 ax1.set_ylim(np.min(self.sys.en.real)-0.2, np.max(self.sys.en.real)+0.2) 

275 ind = np.ones(len(self.sys.en), bool) 

276 else: 

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

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

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

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

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

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

283 for label in ax1.get_yticklabels(): 

284 label.set_color('b') 

285 if tag_pola: 

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

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

288 elif ipr: 

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

290 elif petermann: 

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

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

293 label.set_fontsize(fs) 

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

295 label.set_fontsize(fs) 

296 xa = ax1.get_xaxis() 

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

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

299 fig.set_layout_engine('tight') 

300 plt.draw() 

301 return fig 

302 

303 def polarization( 

304 self, 

305 fig: Figure | None = None, 

306 ax1: Axes | None = None, 

307 ms: float = 10., 

308 fs: float = 20., 

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

310 tag_pola: str | None = None, 

311 ind: NDArray | None = None, 

312 ) -> tuple[Figure, Axes]: 

313 ''' 

314 Plot sublattice polarization. 

315 

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

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

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

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

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

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

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

323 

324 :returns: 

325 * **fig** -- Figure. 

326 ''' 

327 if fig is None: 

328 error_handling.sys(self.sys) 

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

330 error_handling.positive_real(ms, 'ms') 

331 error_handling.positive_real(fs, 'fs') 

332 error_handling.lims(lims) 

333 fig, ax2 = plt.subplots() 

334 if lims is None: 

335 ax2.set_ylim(-0.1, 1.1) 

336 ind = np.ones(len(self.sys.en), bool) 

337 else: 

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

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

340 else: 

341 ax2 = plt.gca().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(len(self.sys.en)) 

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 = '$<' + str(tag_pola).upper() + '|' + str(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 the Inverse Participation Ratio. 

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 if lims is None: 

387 ind = np.ones(len(self.sys.en), bool) 

388 else: 

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

390 else: 

391 ax2 = plt.gca().twinx() 

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

393 x = np.arange(len(self.sys.en)) 

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

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

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

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

398 tick.label1.set_fontsize(fs) 

399 for label in ax2.get_yticklabels(): 

400 label.set_color('r') 

401 return fig, ax2 

402 

403 def petermann( 

404 self, 

405 fig: Figure | None = None, 

406 ax1: Axes | None = None, 

407 ms: float = 10, 

408 fs: float = 20, 

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

410 ind: NDArray | None = None, 

411 ) -> tuple[Figure, Axes]: 

412 ''' 

413 Plot the Petermann factor. 

414 

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

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

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

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

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

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

421 

422 :returns: 

423 * **fig** -- Figure. 

424 ''' 

425 if fig is None: 

426 error_handling.sys(self.sys) 

427 error_handling.empty_ndarray(self.sys.petermann, 'sys.get_petermann') 

428 error_handling.positive_real(ms, 'ms') 

429 error_handling.positive_real(fs, 'fs') 

430 error_handling.lims(lims) 

431 fig, ax2 = plt.subplots() 

432 if lims is None: 

433 ind = np.ones(len(self.sys.en), bool) 

434 else: 

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

436 else: 

437 ax2 = plt.gca().twinx() 

438 error_handling.empty_ndarray(self.sys.petermann, 'sys.get_petermann') 

439 x = np.arange(len(self.sys.en)) 

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

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

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

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

444 tick.label1.set_fontsize(fs) 

445 for label in ax2.get_yticklabels(): 

446 label.set_color('r') 

447 return fig, ax2 

448 

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

450 ''' 

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

452 and imaginary part (red circles). 

453 

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

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

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

457 

458 :returns: 

459 * **fig** -- Figure. 

460 ''' 

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

462 error_handling.positive_real(ms, 'ms') 

463 error_handling.positive_real(fs, 'fs') 

464 error_handling.lims(lims) 

465 fig, ax1 = plt.subplots() 

466 ax1 = plt.gca() 

467 x = np.arange(len(self.sys.en)) 

468 if lims is None: 

469 ax1.set_ylim(min(np.min(self.sys.en.real), np.min(self.sys.en.imag))-0.2, 

470 max(np.max(self.sys.en.real), np.max(self.sys.en.imag))+0.2) 

471 ind = np.ones(len(self.sys.en), bool) 

472 else: 

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

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

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

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

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

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

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

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

481 label.set_fontsize(fs) 

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

483 label.set_fontsize(fs) 

484 xa = ax1.get_xaxis() 

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

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

487 fig.set_layout_engine('tight') 

488 plt.draw() 

489 return fig 

490 

491 def intensity_1d( 

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

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

494 ) -> Figure: 

495 r''' 

496 Plot intensity for 1D lattices. 

497 

498 :param intensity: np.array. Field intensity. 

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

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

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

502 :param title: String. Default value '$|\psi^{(j)}|^2$'. Figure title. 

503 

504 :returns: 

505 * **fig** -- Figure. 

506 ''' 

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

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

509 error_handling.positive_real(ms, 'ms') 

510 error_handling.positive_real(lw, 'lw') 

511 error_handling.positive_real(fs, 'fs') 

512 error_handling.string(title, 'title') 

513 fig, ax = plt.subplots() 

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

515 ax.set_ylabel(title, fontsize=fs) 

516 ax.set_title(title, fontsize=fs) 

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

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

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

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

521 plt.xlim(-1., self.sys.lat.sites) 

522 plt.ylim(0., np.max(intensity)+.05) 

523 fig.set_layout_engine('tight') 

524 plt.draw() 

525 return fig 

526 

527 def intensity_disk( 

528 self, 

529 intensity: NDArray, 

530 s: float = 200., 

531 fs: float = 20., 

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

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

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

535 ) -> Figure: 

536 r''' 

537 Plot the intensity. Colormap with identical disk shape. 

538 

539 :param intensity: np.array.Field intensity. 

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

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

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

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

544 :param title: String. Default value '$|\psi|^2$'. Title. 

545 

546 :returns: 

547 * **fig** -- Figure. 

548 ''' 

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

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

551 error_handling.positive_real(s, 's') 

552 error_handling.positive_real(fs, 'fs') 

553 error_handling.tuple_2elem(figsize, 'figsize') 

554 error_handling.string(title, 'title') 

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

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

557 map_red = plt.get_cmap('Reds') 

558 if lims is None: 

559 lims = (0., np.max(intensity)) 

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

561 else: 

562 y_ticks = [str(v) for v in lims] 

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

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

565 cbar = plt.colorbar(ticks=lims) 

566 ax.set_xticks([]) 

567 ax.set_yticks([]) 

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

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

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

571 cbar.ax.tick_params(labelsize=fs) 

572 ax.set_aspect('equal') 

573 fig.set_layout_engine('tight') 

574 plt.draw() 

575 return fig 

576 

577 def intensity_area( 

578 self, 

579 intensity: NDArray, 

580 s: float = 1000., 

581 lw: float = 1., 

582 fs: float = 20., 

583 plt_hop: bool = False, 

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

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

586 ) -> Figure: 

587 r''' 

588 Plot the intensity. Intensity propotional to disk shape. 

589 

590 :param intensity: np.array. Intensity. 

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

592 Circle size given by 100 * s * intensity. 

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

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

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

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

597 :param title: String. Default value '$|\psi|^2$'. Figure title. 

598 

599 :returns: 

600 * **fig** -- Figure. 

601 ''' 

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

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

604 error_handling.positive_real(s, 's') 

605 error_handling.positive_real(fs, 'fs') 

606 error_handling.boolean(plt_hop, 'plt_hop') 

607 error_handling.tuple_2elem(figsize, 'figsize') 

608 error_handling.string(title, 'title') 

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

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

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

612 ax.set_title(title, fontsize=fs) 

613 if plt_hop: 

614 hop = self.sys.hop[~self.sys._wrapped(self.sys.hop['i'].astype(int), 

615 self.sys.hop['j'].astype(int))] 

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

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

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

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

620 'k', lw=lw) 

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

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

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

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

625 c=color, alpha=0.5) 

626 ax.set_aspect('equal') 

627 ax.axis('off') 

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

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

630 ax.set_xlim(x_lim) 

631 ax.set_ylim(y_lim) 

632 fig.set_layout_engine('tight') 

633 plt.draw() 

634 return fig 

635 

636 def butterfly( 

637 self, 

638 betas: NDArray, 

639 butterfly: NDArray, 

640 lw: float = 1., 

641 fs: float = 20., 

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

643 title: str = '', 

644 ) -> Figure: 

645 ''' 

646 Plot energies depending on the strain (see 

647 *GrapheneSystem.get_butterfly*). The energy levels plotted are the 

648 ones lying in *lims* at the smallest strain. 

649 

650 :param betas: np.array. Strain values, shape (N,).