Coverage for tbkit/propagation.py: 100%

1from __future__ import annotations

3import numpy as np

4from numpy.typing import NDArray

5import scipy.sparse as sparse

6import scipy.linalg as LA

7import matplotlib.pyplot as plt

8from matplotlib.figure import Figure

9from matplotlib.animation import FuncAnimation

10import os

11import tbkit.error_handling as error_handling

12from tbkit.lattice import Lattice

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

17# CLASS PROPAGATION

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

21class Propagation():

22 r'''

23 Get lattice time evolution. The time-dependent Schrodinger equation

24 :math:`i\,\partial_z\psi = H\psi` is solved by the Crank-Nicolson method,

26 .. math::

28 \psi(z+dz) = \left(1 + \tfrac{i}{2}H\,dz\right)^{-1}

29 \left(1 - \tfrac{i}{2}H\,dz\right)\psi(z)\, ,

31 unitary for a Hermitian :math:`H`.

33 Parameters

34 ----------

35 lat : **lattice** class instance

36 '''

38 def __init__(self, lat: Lattice) -> None:

39 error_handling.lat(lat)

40 self.lat = lat

41 self.prop = np.array([], 'c16')

43 def get_propagation(

44 self, ham: sparse.spmatrix, psi_init: NDArray, steps: int, dz: float, norm: bool = False,

45 ) -> None:

46 r'''

47 Get the time evolution.

49 Parameters

50 ----------

51 ham : sparse.csr_matrix

52 Tight-Binding Hamiltonian.

53 psi_init : np.ndarray

54 Initial state.

55 steps : Positive Integer

56 Number of steps.

57 dz : Positive number

58 Step.

59 norm : Boolean, default False

60 Renormalize the state to

61 :math:`\sum_i|\psi_i|^2 = 1` at each step. The Crank-Nicolson

62 step is unitary, so this only matters for a non-Hermitian

63 Hamiltonian (gain/loss), where the norm is not conserved.

64 '''

65 error_handling.empty_ham(ham)

66 error_handling.ndarray(psi_init, 'psi_init', self.lat.sites)

67 error_handling.positive_int(steps, 'steps')

68 error_handling.positive_real(dz, 'dz')

69 error_handling.boolean(norm, 'norm')

70 self.steps = steps

71 self.dz = dz

72 self.prop = np.empty((self.lat.sites, self.steps), 'c16')

73 self.prop[:, 0] = psi_init

74 diag = 1j*np.ones(self.lat.sites, 'c16')

75 A = (sparse.diags(diag, 0) - 0.5 * self.dz * ham).toarray()

76 B = (sparse.diags(diag, 0) + 0.5 * self.dz * ham).toarray()

77 mat = (np.dot(LA.inv(A), B))

78 for i in range(1, self.steps):

79 self.prop[:, i] = np.dot(mat, self.prop[:, i-1])

80 if norm:

81 self.prop[:, i] /= np.linalg.norm(self.prop[:, i])

83 def get_pumping(

84 self, hams: list[sparse.spmatrix], psi_init: NDArray, steps: int, dz: float, norm: bool = True,

85 ) -> None:

86 r'''

87 Get the time evolution with adiabatic pumpings.

89 Parameters

90 ----------

91 hams : List of sparse.csr_matrices

92 Tight-Binding Hamiltonians.

93 psi_init : np.ndarray

94 Initial state.

95 steps : Positive integer

96 Number of steps.

97 dz : Positive number

98 Step.

99 norm : Boolean, default True

100 Renormalize the state to

101 :math:`\sum_i|\psi_i|^2 = 1` at each step.

102

103 Notes

104 -----

105 The *steps* are split into ``len(hams) + 1`` equal stages: evolution

106 under ``hams[0]``, then ``len(hams) - 1`` stages each interpolating

107 linearly from ``hams[j]`` to ``hams[j+1]``, then evolution under

108 ``hams[-1]`` for the remaining steps. With fewer than

109 ``len(hams) + 1`` steps the stages are empty: the state evolves under

110 ``hams[-1]`` from the first step on.

111 '''

112 error_handling.get_pump(hams)

113 error_handling.ndarray(psi_init, 'psi_init', self.lat.sites)

114 error_handling.positive_int(steps, 'steps')

115 error_handling.positive_real(dz, 'dz')

116 error_handling.boolean(norm, 'norm')

117 self.steps = steps

118 self.dz = dz

119 no = len(hams)

120 self.prop = np.empty((self.lat.sites, self.steps), 'c16')

121 self.prop[:, 0] = psi_init

122 diag = 1j * np.ones(self.lat.sites, 'c16')

123 delta = self.steps // (1 + no)

124 A = (sparse.diags(diag, 0) - 0.5 * self.dz * hams[0]).toarray()

125 B = (sparse.diags(diag, 0) + 0.5 * self.dz * hams[0]).toarray()

126 mat = (np.dot(LA.inv(A), B))

127 # before pumping

128 for i in range(1, delta):

129 self.prop[:, i] = np.dot(mat, self.prop[:, i-1])

130 if norm:

131 self.prop[:, i] /= np.linalg.norm(self.prop[:, i])

132 # pumping

133 c = np.linspace(0, 1, delta)

134 for j in range(0, no-1):

135 for i in range(0, delta):

136 ham = (1-c[i])*hams[j]+c[i]*hams[j+1]

137 A = (sparse.diags(diag, 0) - 0.5 * self.dz * ham).toarray()

138 B = (sparse.diags(diag, 0) + 0.5 * self.dz * ham).toarray()

139 mat = (np.dot(LA.inv(A), B))

140 self.prop[:, (j+1)*delta+i] = np.dot(mat, self.prop[:, (j+1)*delta+i-1])

141 if norm:

142 self.prop[:, (j+1)*delta+i] /= \

143 np.linalg.norm(self.prop[:, (j+1)*delta+i])

144 # after pumping, under hams[-1] (what the last pumping step ends on).

145 # Start at 1 at least: with delta == 0 there were no stages, and

146 # index no*delta - 1 = -1 would overwrite the initial state.

147 A = (sparse.diags(diag, 0) - 0.5 * self.dz * hams[-1]).toarray()

148 B = (sparse.diags(diag, 0) + 0.5 * self.dz * hams[-1]).toarray()

149 mat = (np.dot(LA.inv(A), B))

150 for i in range(max(no*delta, 1), self.steps):

151 self.prop[:, i] = np.dot(mat, self.prop[:, i-1])

152 if norm:

153 self.prop[:, i] /= np.linalg.norm(self.prop[:, i])

154

155 def plt_propagation_1d(

156 self, prop_type: str = 'real', fs: float = 20, figsize: tuple[float, float] | None = None,

157 ) -> Figure:

158 '''

159 Plot time evolution for 1D systems.

160

161 Parameters

162 ----------

163 prop_type : String, default 'real'

164 Plot the real part

165 ('real'), the imaginary part ('imag'), or the intensity ('norm').

166 fs : Positive number, default 20

167 Fontsize.

168 figsize : Tuple, default None

169 Figure size.

170

171 Returns

172 -------

173 fig : Figure

174 '''

175 error_handling.empty_ndarray(self.prop, 'get_propagation or get_pumping')

176 error_handling.positive_real(fs, 'fs')

177 error_handling.prop_type(prop_type)

178 error_handling.tuple_2elem(figsize, 'figsize')

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

180 plt.ylabel('n', fontsize=fs)

181 plt.xlabel('z', fontsize=fs)

182 if prop_type == 'real':

183 color = self.prop_smooth_1d(self.prop.real)

184 max_val = max(np.max(color), -np.min(color))

185 ticks = [-max_val, max_val]

186 cmap = 'seismic'

187 elif prop_type == 'imag':

188 color = self.prop_smooth_1d(self.prop.imag)

189 max_val = max(np.max(color), -np.min(color))

190 ticks = [-max_val, max_val]

191 cmap = 'seismic'

192 else:

193 color = self.prop_smooth_1d(np.abs(self.prop) ** 2)

194 ticks = [0., np.max(color[:, -1])]

195 cmap = 'hot'

196 extent = (-0, self.steps*self.dz, self.lat.sites-.5, -.5)

197 interpolation = 'nearest'

198 im = plt.imshow(color, cmap=cmap, aspect='auto',

199 interpolation=interpolation, extent=extent,

200 vmin=ticks[0], vmax=ticks[-1])

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

202 label.set_fontsize(fs)

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

204 label.set_fontsize(fs)

205 ax.get_yaxis().set_major_locator(plt.MaxNLocator(integer=True))

206 if prop_type == 'norm':

207 cbar = fig.colorbar(im, ticks=ticks)

208 cbar.ax.set_yticklabels(['0','max'])

209 else:

210 cbar = fig.colorbar(im, ticks=[ticks[0], 0, ticks[1]])

211 cbar.ax.set_yticklabels(['min', '0','max'])

212 cbar.ax.tick_params(labelsize=fs)

213 return fig

214

215 def prop_smooth_1d(self, prop: NDArray, a: float = 10, no: int = 40) -> NDArray:

216 r'''

217 Private function. Used in *plt_propagation_1d*.

218 Smooth propagation for 1D systems.

219 Perform Gaussian interpolation :math:`e^{-a(x-x_i)^2}`,

220

221 Parameters

222 ----------

223 prop

224 Propagation.

225 a

226 Default value 10. Gaussian Parameter.

227 no

228 Default value 40. Number of points of each Gaussian.

229

230 Returns

231 -------

232 smooth

233 Smoothed propagation.

234 '''

235 func = np.exp(- a * np.linspace(-0.5, 0.5, no) ** 2)

236 smooth = np.empty((self.lat.sites * no, self.steps))

237 for iz in range(0, self.steps):

238 for i in range(self.lat.sites):

239 smooth[i*no: (i+1)*no, iz] = prop[i, iz] * func

240 return smooth

241

242 def get_animation(

243 self, s: float = 300., fs: float = 20., prop_type: str = 'real',

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

245 ) -> FuncAnimation:

246 '''

247 Get time evolution animation.

248

249 Parameters

250 ----------

251 s

252 Default value 300. Circle size.

253 fs

254 Default value 20. Fontsize.

255 prop_type : String, default 'real'

256 Plot the real part

257 ('real'), the imaginary part ('imag'), or the intensity ('norm').

258 figsize : Tuple, default None

259 Figsize.

260

261 Returns

262 -------

263 ani

264 Animation.

265 '''

266 error_handling.empty_ndarray(self.prop, 'get_propagation or get_pumping')

267 error_handling.positive_real(s, 's')

268 error_handling.positive_real(fs, 'fs')

269 error_handling.prop_type(prop_type)

270 error_handling.tuple_2elem(figsize, 'figsize')

271 if os.name == 'posix':

272 blit = False

273 else:

274 blit = True

275 if prop_type == 'real':

276 color = self.prop.real

277 max_val = max(np.max(color), -np.min(color))

278 ticks = [-max_val, max_val]

279 cmap = 'seismic'

280 elif prop_type == 'imag':

281 color = self.prop.imag

282 max_val = max(np.max(color), -np.min(color))

283 ticks = [-max_val, max_val]

284 cmap = 'seismic'

285 else:

286 color = np.abs(self.prop) ** 2

287 ticks = [0., np.max(color)]

288 cmap = 'Reds'

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

290 plt.xlim(np.min(self.lat.coor['x'])-1., np.max(self.lat.coor['x'])+1.)

291 plt.ylim(np.min(self.lat.coor['y'])-1., np.max(self.lat.coor['y'])+1.)

292 scat = plt.scatter(self.lat.coor['x'], self.lat.coor['y'], c=color[:, 0],

293 s=s, vmin=ticks[0], vmax=ticks[1],

294 cmap=plt.get_cmap(cmap))

295 frame = plt.gca()

296 frame.get_xaxis().set_ticks([])

297 frame.get_yaxis().set_ticks([])

298 ax.set_aspect('equal')

299 if prop_type == 'norm':

300 cbar = fig.colorbar(scat, ticks=ticks)

301 cbar.ax.set_yticklabels(['0','max'])

302 else:

303 cbar = fig.colorbar(scat, ticks=[ticks[0], 0, ticks[1]])

304 cbar.ax.set_yticklabels(['min', '0','max'])

305

306 def update(i, color, scat):

307 scat.set_array(color[:, i])

308 return scat,

309

310 ani = FuncAnimation(fig, update, frames=self.steps,

311 fargs=(color, scat), blit=blit, repeat=False)

312 return ani

313

314 def get_animation_nb(

315 self, s: float = 300., fs: float = 20., prop_type: str = 'real',

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

317 ) -> FuncAnimation:

318 '''

319 Get time evolution animation for iPython notebooks.

320

321 Parameters

322 ----------

323 s

324 Default value 300. Circle size.

325 fs

326 Default value 20. Fontsize.

327 prop_type : String, default 'real'

328 Plot the real part

329 ('real'), the imaginary part ('imag'), or the intensity ('norm').

330 figsize : Tuple, default None

331 Figsize.

332

333 Returns

334 -------

335 ani

336 Animation.

337 '''

338 error_handling.empty_ndarray(self.prop, 'get_propagation or get_pumping')

339 error_handling.positive_real(s, 's')

340 error_handling.positive_real(fs, 'fs')

341 error_handling.prop_type(prop_type)

342 error_handling.tuple_2elem(figsize, 'figsize')

343 if prop_type == 'real' or prop_type == 'imag':

344 color = self.prop.real if prop_type == 'real' else self.prop.imag

345 max_val = max(np.max(color), -np.min(color))

346 ticks = [-max_val, max_val]

347 cmap = 'seismic'

348 else:

349 color = np.abs(self.prop) ** 2

350 ticks = [0., np.max(color)]

351 cmap = 'Reds'

352 fig = plt.figure(figsize=figsize)

353 ax = plt.axes(xlim=(np.min(self.lat.coor['x']-.5), np.max(self.lat.coor['x']+.5)),

354 ylim=(np.min(self.lat.coor['y']-.5), np.max(self.lat.coor['y']+.5)))

355 ax.set_aspect('equal')

356 frame = plt.gca()

357 frame.get_xaxis().set_ticks([])

358 frame.get_yaxis().set_ticks([])

359 scat = plt.scatter(self.lat.coor['x'], self.lat.coor['y'], c=color[:, 0],

360 s=s, vmin=ticks[0], vmax=ticks[1],

361 cmap=cmap)

362 if prop_type == 'real' or prop_type == 'imag':

363 cbar = fig.colorbar(scat, ticks=[ticks[0], 0, ticks[1]])

364 cbar.ax.set_yticklabels(['min', '0','max'])

365 else:

366 cbar = fig.colorbar(scat, ticks=[0, ticks[1]])

367 cbar.ax.set_yticklabels(['0','max'])

368

369 def init():

370 scat.set_array(color[:, 0])

371 return scat,

372

373 def animate(i):

374 scat.set_array(color[:, i])

375 return scat,

376

377 return FuncAnimation(fig, animate, init_func=init,

378 frames=self.steps, interval=120, blit=True)

379

380 def plt_prop_dimer(self, lw: float = 5, fs: float = 20) -> Figure:

381 '''

382 Plot time evolution for dimers.

383

384 Parameters

385 ----------

386 lw : Positive number, default 5

387 Linewidth.

388 fs : Positive number, default 20

389 Fontsize.

390

391 Returns

392 -------

393 fig : Figure

394 '''

395 error_handling.empty_ndarray(self.prop, 'get_propagation or get_pumping')

396 error_handling.positive_real(lw, 'lw')

397 error_handling.positive_real(fs, 'fs')

398 color = ['b', 'r']

399 fig, ax = plt.subplots()

400 z = self.dz * np.arange(self.steps)

401 for i, c in zip([0, 1], color):

402 plt.plot(z, np.abs(self.prop[i, :])**2, c, lw=lw)

403 plt.title('Intensity', fontsize=fs)

404 plt.xlabel('$z$', fontsize=fs)

405 plt.ylabel(r'$|\psi_j|^2$', fontsize=fs)

406 plt.xlim(0, z[-1])

407 return fig

408

409

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

411propagation = Propagation