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 Parameters 

23 ---------- 

24 sys : class instance **system** 

25 colors 

26 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, ax: Axes | None = None) -> None: 

38 ''' 

39 Private method called by *lattice_generic*. 

40 ''' 

41 if ax is None: 

42 ax = plt.gca() 

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

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

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

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

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

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

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

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

51 

52 def lattice_generic( 

53 self, 

54 coor: NDArray, 

55 ms: float, 

56 lw: float, 

57 c: float, 

58 fs: float, 

59 axis: bool, 

60 plt_hop: bool, 

61 plt_hop_low: bool, 

62 plt_index: bool, 

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

64 ax: Axes | None = None, 

65 ) -> Figure | SubFigure: 

66 ''' 

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

68 ''' 

69 error_handling.positive_real(ms, 'ms') 

70 error_handling.positive_real(lw, 'lw') 

71 error_handling.positive_real(c, 'c') 

72 error_handling.positive_real(fs, 'fs') 

73 error_handling.boolean(axis, 'axis') 

74 error_handling.boolean(plt_hop, 'plt_hop') 

75 error_handling.boolean(plt_hop_low, 'plt_hop_low') 

76 error_handling.boolean(plt_index, 'plt_index') 

77 error_handling.tuple_2elem(figsize, 'figsize') 

78 fig: Figure | SubFigure 

79 if ax is None: 

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

81 else: 

82 fig = ax.figure 

83 # hoppings 

84 if plt_hop: 

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 if plt_hop_low: 

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

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

90 # plot sites 

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

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

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

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

95 ax.set_aspect('equal') 

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

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

98 if not axis: 

99 ax.axis('off') 

100 # plot indices 

101 if plt_index: 

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

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

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

105 textcoords='offset points', 

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

107 plt.draw() 

108 return fig 

109 

110 def lattice( 

111 self, 

112 ms: float = 20, 

113 lw: float = 5., 

114 c: float = 3., 

115 fs: float = 20, 

116 axis: bool = False, 

117 plt_hop: bool = False, 

118 plt_hop_low: bool = False, 

119 plt_index: bool = False, 

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

121 ax: Axes | None = None, 

122 ) -> Figure | SubFigure: 

123 ''' 

124 Plot lattice. 

125 

126 Parameters 

127 ---------- 

128 ms : Positive number, default 20 

129 Markersize. 

130 c : Positive number, default 3 

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

132 fs : Positive number, default 20 

133 Fontsize. 

134 plt_hop : Boolean, default False 

135 Plot hoppings. 

136 plt_hop_low : Boolean, default False 

137 Plot hoppings diagonal low. 

138 plt_index : Boolean, default False 

139 Plot site labels. 

140 axis : Boolean, default False 

141 Plot axis. 

142 figsize : Tuple, default None 

143 Figure size. 

144 ax : Axes, default None 

145 Draw on this axis instead of 

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

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

148 

149 Returns 

150 ------- 

151 fig : Figure 

152 ''' 

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

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

155 plt_hop_low, plt_index, figsize, ax) 

156 

157 def lattice_hop( 

158 self, 

159 ms: float = 20, 

160 lw: float = 5, 

161 c: float = 3., 

162 fs: float = 20, 

163 axis: bool = False, 

164 plt_hop: bool = False, 

165 plt_hop_low: bool = False, 

166 plt_index: bool = False, 

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

168 ax: Axes | None = None, 

169 ) -> Figure | SubFigure: 

170 ''' 

171 Plot lattice in hopping space. 

172 

173 Parameters 

174 ---------- 

175 ms : Positive Float, default 20 

176 Markersize. 

177 c : Positive Float, default 3 

178 Coefficient. 

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

180 fs : Positive Float, default 20 

181 Fontsize. 

182 axis : Boolean, default False 

183 Plot axis. 

184 plt_hop : Boolean, default False 

185 Plot hoppings. 

186 plt_index : Boolean, default False 

187 Plot site labels. 

188 figsize : Tuple, default None 

189 Figure size. 

190 ax : Axes, default None 

191 Draw on this axis instead of 

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

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

194 

195 Returns 

196 ------- 

197 fig : Figure 

198 ''' 

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

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

201 plt_hop_low, plt_index, figsize, ax) 

202 

203 

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

205 """ 

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

207 

208 Parameters 

209 ---------- 

210 nbr_bins : Positive integer, default 61 

211 Number of bins of the histogram. 

212 fs : Positive number, default 20 

213 Fontsize. 

214 lims : List, lims[0] energy min, lims[1] energy max 

215 

216 Returns 

217 ------- 

218 fig : Figure 

219 """ 

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

221 error_handling.positive_int(nbr_bins, 'nbr_bins') 

222 error_handling.positive_real(fs, 'fs') 

223 error_handling.lims(lims) 

224 fig, ax = plt.subplots() 

225 en_real = self.sys.en.real 

226 if lims is None: 

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

228 else: 

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

230 ax.set_xlim(lims) 

231 en = en_real[ind_en] 

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

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

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

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

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

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

238 label.set_fontsize(fs) 

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

240 label.set_fontsize(fs) 

241 return fig 

242 

243 def dos( 

244 self, 

245 broadening: float = 0.05, 

246 kernel: str = 'gaussian', 

247 e_grid: ArrayLike | None = None, 

248 fs: float = 20, 

249 lw: float = 2., 

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

251 ) -> Figure: 

252 ''' 

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

254 

255 Parameters 

256 ---------- 

257 broadening : Positive real number, default 0.05 

258 Kernel width. 

259 kernel : String, default 'gaussian' 

260 'gaussian' or 'lorentzian'. 

261 e_grid : Real ndarray, default None 

262 Energies at which to 

263 evaluate the density of states. 

264 fs : Positive number, default 20 

265 Fontsize. 

266 lw : Positive number, default 2 

267 Linewidth. 

268 figsize : Tuple, default None 

269 Figure size. 

270 

271 Returns 

272 ------- 

273 fig : Figure 

274 ''' 

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

276 error_handling.positive_real(fs, 'fs') 

277 error_handling.positive_real(lw, 'lw') 

278 error_handling.tuple_2elem(figsize, 'figsize') 

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

280 

281 def spectrum( 

282 self, 

283 ms: float = 10, 

284 fs: float = 20, 

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

286 tag_pola: str | None = None, 

287 ipr: bool | None = None, 

288 petermann: bool | None = None, 

289 ) -> Figure: 

290 ''' 

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

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

293 

294 Parameters 

295 ---------- 

296 ms 

297 Default value 10. Markersize. 

298 fs 

299 Default value 20. Fontsize. 

300 lims : List, lims[0] energy min, lims[1] energy max 

301 tag_pola 

302 Default value None. One-character string. Tag of the sublattice. 

303 ipr 

304 Default value None. If True plot the Inverse Participation Ratio. 

305 petermann 

306 Default value None. If True plot the Petermann factor. 

307 

308 Returns 

309 ------- 

310 fig : Figure 

311 ''' 

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

313 error_handling.positive_real(ms, 'ms') 

314 error_handling.positive_real(fs, 'fs') 

315 error_handling.lims(lims) 

316 fig, ax1 = plt.subplots() 

317 ax1 = plt.gca() 

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

319 if lims is None: 

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

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

322 else: 

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

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

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

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

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

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

329 for label in ax1.get_yticklabels(): 

330 label.set_color('b') 

331 if tag_pola: 

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

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

334 elif ipr: 

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

336 elif petermann: 

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

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

339 label.set_fontsize(fs) 

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

341 label.set_fontsize(fs) 

342 xa = ax1.get_xaxis() 

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

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

345 fig.set_layout_engine('tight') 

346 plt.draw() 

347 return fig 

348 

349 def polarization( 

350 self, 

351 fig: Figure | None = None, 

352 ax1: Axes | None = None, 

353 ms: float = 10., 

354 fs: float = 20., 

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

356 tag_pola: str | None = None, 

357 ind: NDArray | None = None, 

358 ) -> tuple[Figure, Axes]: 

359 ''' 

360 Plot sublattice polarization. 

361 

362 Parameters 

363 ---------- 

364 fig : Figure, default None 

365 (used by the method spectrum). 

366 ax1 : Axis, default None 

367 (used by the method spectrum). 

368 ms : Positive Float, default 10 

369 Markersize. 

370 fs : Positive Float, default 20 

371 Fontsize. 

372 lims : List, lims[0] energy min, lims[1] energy max 

373 tag_pola : One-character string, default None 

374 Tag of the sublattice. 

375 ind : List, default None 

376 List of indices. (used in the method spectrum). 

377 

378 Returns 

379 ------- 

380 fig : Figure 

381 ''' 

382 if fig is None: 

383 error_handling.sys(self.sys) 

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

385 error_handling.positive_real(ms, 'ms') 

386 error_handling.positive_real(fs, 'fs') 

387 error_handling.lims(lims) 

388 fig, ax2 = plt.subplots() 

389 if lims is None: 

390 ax2.set_ylim(-0.1, 1.1) 

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

392 else: 

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

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

395 else: 

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

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

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

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

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

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

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

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

404 ax2.set_ylim(-0.1, 1.1) 

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

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

407 tick.label1.set_fontsize(fs) 

408 for label in ax2.get_yticklabels(): 

409 label.set_color('r') 

410 return fig, ax2 

411 

412 def ipr( 

413 self, 

414 fig: Figure | None = None, 

415 ax1: Axes | None = None, 

416 ms: float = 10, 

417 fs: float = 20, 

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

419 ind: NDArray | None = None, 

420 ) -> tuple[Figure, Axes]: 

421 ''' 

422 Plot the Inverse Participation Ratio. 

423 

424 Parameters 

425 ---------- 

426 fig : Figure, default None 

427 (used by the method spectrum). 

428 ax1 : Axis, default None 

429 (used by the method spectrum). 

430 ms : Positive Float, default 10 

431 Markersize. 

432 fs : Positive Float, default 20 

433 Fontsize. 

434 lims : List 

435 lims[0] energy min, lims[1] energy max. 

436 ind : List, default None 

437 List of indices. (used in the method spectrum). 

438 

439 Returns 

440 ------- 

441 fig : Figure 

442 ''' 

443 if fig is None: 

444 error_handling.sys(self.sys) 

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

446 error_handling.positive_real(ms, 'ms') 

447 error_handling.positive_real(fs, 'fs') 

448 error_handling.lims(lims) 

449 fig, ax2 = plt.subplots() 

450 if lims is None: 

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

452 else: 

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

454 else: 

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

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

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

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

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

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

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

462 tick.label1.set_fontsize(fs) 

463 for label in ax2.get_yticklabels(): 

464 label.set_color('r') 

465 return fig, ax2 

466 

467 def petermann( 

468 self, 

469 fig: Figure | None = None, 

470 ax1: Axes | None = None, 

471 ms: float = 10, 

472 fs: float = 20, 

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

474 ind: NDArray | None = None, 

475 ) -> tuple[Figure, Axes]: 

476 ''' 

477 Plot the Petermann factor. 

478 

479 Parameters 

480 ---------- 

481 fig : Figure, default None 

482 (used by the method spectrum). 

483 ax1 : Axis, default None 

484 (used by the method spectrum). 

485 ms : Positive Float, default 10 

486 Markersize. 

487 fs : Positive Float, default 20 

488 Fontsize. 

489 lims : List 

490 lims[0] energy min, lims[1] energy max. 

491 ind : List, default None 

492 List of indices. (used in the method spectrum). 

493 

494 Returns 

495 ------- 

496 fig : Figure 

497 ''' 

498 if fig is None: 

499 error_handling.sys(self.sys) 

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

501 error_handling.positive_real(ms, 'ms') 

502 error_handling.positive_real(fs, 'fs') 

503 error_handling.lims(lims) 

504 fig, ax2 = plt.subplots() 

505 if lims is None: 

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

507 else: 

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

509 else: 

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

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

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

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

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

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

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

517 tick.label1.set_fontsize(fs) 

518 for label in ax2.get_yticklabels(): 

519 label.set_color('r') 

520 return fig, ax2 

521 

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

523 ''' 

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

525 and imaginary part (red circles). 

526 

527 Parameters 

528 ---------- 

529 ms : Positive Float, default 10 

530 Markersize. 

531 fs : Positive Float, default 20 

532 Font size. 

533 lims : List 

534 lims[0] energy min, lims[1] energy max. 

535 

536 Returns 

537 ------- 

538 fig : Figure 

539 ''' 

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

541 error_handling.positive_real(ms, 'ms') 

542 error_handling.positive_real(fs, 'fs') 

543 error_handling.lims(lims) 

544 fig, ax1 = plt.subplots() 

545 ax1 = plt.gca() 

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

547 if lims is None: 

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

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

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

551 else: 

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

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

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

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

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

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

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

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

560 label.set_fontsize(fs) 

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

562 label.set_fontsize(fs) 

563 xa = ax1.get_xaxis() 

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

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

566 fig.set_layout_engine('tight') 

567 plt.draw() 

568 return fig 

569 

570 def intensity_1d( 

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

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

573 ) -> Figure: 

574 r''' 

575 Plot intensity for 1D lattices. 

576 

577 Parameters 

578 ---------- 

579 intensity : np.array 

580 Field intensity. 

581 ms : Positive Float, default 20 

582 Markersize. 

583 lw : Positive Float, default 2 

584 Linewith, connect sublattice sites. 

585 fs : Positive Float, default 20 

586 Font size. 

587 title : String 

588 Default value '$|\psi^{(j)}|^2$'. Figure title. 

589 

590 Returns 

591 ------- 

592 fig : Figure 

593 ''' 

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

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

596 error_handling.positive_real(ms, 'ms') 

597 error_handling.positive_real(lw, 'lw') 

598 error_handling.positive_real(fs, 'fs') 

599 error_handling.string(title, 'title') 

600 fig, ax = plt.subplots() 

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

602 ax.set_ylabel(title, fontsize=fs) 

603 ax.set_title(title, fontsize=fs) 

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

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

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

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

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

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

610 fig.set_layout_engine('tight') 

611 plt.draw() 

612 return fig 

613 

614 def intensity_disk( 

615 self, 

616 intensity: NDArray, 

617 s: float = 200., 

618 fs: float = 20., 

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

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

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

622 ) -> Figure: 

623 r''' 

624 Plot the intensity. Colormap with identical disk shape. 

625 

626 Parameters 

627 ---------- 

628 intensity : np.array 

629 Field intensity. 

630 s 

631 Default value 200. Disk size. 

632 fs 

633 Default value 20. Font size. 

634 lims : List, default None 

635 Colormap limits. 

636 figsize : Tuple, default None 

637 Figure size. 

638 title : String, default '$|\psi|^2$' 

639 Title. 

640 

641 Returns 

642 ------- 

643 fig : Figure 

644 ''' 

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

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

647 error_handling.positive_real(s, 's') 

648 error_handling.positive_real(fs, 'fs') 

649 error_handling.tuple_2elem(figsize, 'figsize') 

650 error_handling.string(title, 'title') 

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

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

653 map_red = plt.get_cmap('Reds') 

654 if lims is None: 

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

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

657 else: 

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

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

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

661 cbar = plt.colorbar(ticks=lims) 

662 ax.set_xticks([]) 

663 ax.set_yticks([]) 

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

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

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

667 cbar.ax.tick_params(labelsize=fs) 

668 ax.set_aspect('equal') 

669 fig.set_layout_engine('tight') 

670 plt.draw() 

671 return fig 

672 

673 def intensity_area( 

674 self, 

675 intensity: NDArray, 

676 s: float = 1000., 

677 lw: float = 1., 

678 fs: float = 20., 

679 plt_hop: bool = False, 

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

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

682 ) -> Figure: 

683 r''' 

684 Plot the intensity. Intensity propotional to disk shape. 

685 

686 Parameters 

687 ---------- 

688 intensity : np.array 

689 Intensity. 

690 s : Positive Float, default 1000 

691 Circle size given by 100 * s * intensity. 

692 lw : Positive Float, default 1 

693 Hopping linewidths. 

694 fs : Positive Float, default 20 

695 Fontsize. 

696 plt_hop : Boolean, default False 

697 Plot hoppings. 

698 figsize : Tuple, default None 

699 Figure size. 

700 title : String, default '$|\psi|^2$' 

701 Figure title. 

702 

703 Returns 

704 ------- 

705 fig : Figure 

706 ''' 

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

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

709 error_handling.positive_real(s, 's') 

710 error_handling.positive_real(fs, 'fs') 

711 error_handling.boolean(plt_hop, 'plt_hop') 

712 error_handling.tuple_2elem(figsize, 'figsize') 

713 error_handling.string(title, 'title') 

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

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