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 :param lat: **lattice** class instance.

34 '''

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

37 error_handling.lat(lat)

38 self.lat = lat

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

41 def get_propagation(

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

43 ) -> None:

44 r'''

45 Get the time evolution.

47 :param ham: sparse.csr_matrix. Tight-Binding Hamiltonian.

48 :param psi_init: np.ndarray. Initial state.

49 :param steps: Positive Integer. Number of steps.

50 :param dz: Positive number. Step.

51 :param norm: Boolean. Default value False. Renormalize the state to

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

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

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

55 '''

56 error_handling.empty_ham(ham)

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

58 error_handling.positive_int(steps, 'steps')

59 error_handling.positive_real(dz, 'dz')

60 error_handling.boolean(norm, 'norm')

61 self.steps = steps

62 self.dz = dz

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

64 self.prop[:, 0] = psi_init

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

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

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

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

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

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

71 if norm:

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

74 def get_pumping(

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

76 ) -> None:

77 r'''

78 Get the time evolution with adiabatic pumpings.

80 :param hams: List of sparse.csr_matrices. Tight-Binding Hamiltonians.

81 :param psi_init: np.ndarray. Initial state.

82 :param steps: Positive integer. Number of steps.

83 :param dz: Positive number. Step.

84 :param norm: Boolean. Default value True. Renormalize the state to

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

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

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

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

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

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

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

93 '''

94 error_handling.get_pump(hams)

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

96 error_handling.positive_int(steps, 'steps')

97 error_handling.positive_real(dz, 'dz')

98 error_handling.boolean(norm, 'norm')

99 self.steps = steps

100 self.dz = dz

101 no = len(hams)

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

103 self.prop[:, 0] = psi_init

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

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

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

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

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

109 # before pumping

110 for i in range(1, delta):

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

112 if norm:

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

114 # pumping

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

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

117 for i in range(0, delta):

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

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

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

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

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

123 if norm:

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

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

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

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

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

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

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

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

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

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

134 if norm:

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

136

137 def plt_propagation_1d(

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

139 ) -> Figure:

140 '''

141 Plot time evolution for 1D systems.

142

143 :param prop_type: String. Default value 'real'. Plot the real part

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

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

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

147

148 :returns:

149 * **fig** -- Figure.

150 '''

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

152 error_handling.positive_real(fs, 'fs')

153 error_handling.prop_type(prop_type)

154 error_handling.tuple_2elem(figsize, 'figsize')

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

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

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

158 if prop_type == 'real':

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

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

161 ticks = [-max_val, max_val]

162 cmap = 'seismic'

163 elif prop_type == 'imag':

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

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

166 ticks = [-max_val, max_val]

167 cmap = 'seismic'

168 else:

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

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

171 cmap = 'hot'

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

173 interpolation = 'nearest'

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

175 interpolation=interpolation, extent=extent,

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

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

178 label.set_fontsize(fs)

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

180 label.set_fontsize(fs)

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

182 if prop_type == 'norm':

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

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

185 else:

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

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

188 cbar.ax.tick_params(labelsize=fs)

189 return fig

190

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

192 r'''

193 Private function. Used in *plt_propagation_1d*.

194 Smooth propagation for 1D systems.

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

196

197 :param prop: Propagation.

198 :param a: Default value 10. Gaussian Parameter.

199 :param no: Default value 40. Number of points of each Gaussian.

200

201 :returns:

202 * **smooth** -- Smoothed propagation.

203 '''

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

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

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

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

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

209 return smooth

210

211 def get_animation(

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

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

214 ) -> FuncAnimation:

215 '''

216 Get time evolution animation.

217

218 :param s: Default value 300. Circle size.

219 :param fs: Default value 20. Fontsize.

220 :param prop_type: String. Default value 'real'. Plot the real part

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

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

223

224 :returns:

225 * **ani** -- Animation.

226 '''

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

228 error_handling.positive_real(s, 's')

229 error_handling.positive_real(fs, 'fs')

230 error_handling.prop_type(prop_type)

231 error_handling.tuple_2elem(figsize, 'figsize')

232 if os.name == 'posix':

233 blit = False

234 else:

235 blit = True

236 if prop_type == 'real':

237 color = self.prop.real

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

239 ticks = [-max_val, max_val]

240 cmap = 'seismic'

241 elif prop_type == 'imag':

242 color = self.prop.imag

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

244 ticks = [-max_val, max_val]

245 cmap = 'seismic'

246 else:

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

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

249 cmap = 'Reds'

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

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

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

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

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

255 cmap=plt.get_cmap(cmap))

256 frame = plt.gca()

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

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

259 ax.set_aspect('equal')

260 if prop_type == 'norm':

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

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

263 else:

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

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

266

267 def update(i, color, scat):

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

269 return scat,

270

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

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

273 return ani

274

275 def get_animation_nb(

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

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

278 ) -> FuncAnimation:

279 '''

280 Get time evolution animation for iPython notebooks.

281

282 :param s: Default value 300. Circle size.

283 :param fs: Default value 20. Fontsize.

284 :param prop_type: String. Default value 'real'. Plot the real part

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

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

287

288 :returns:

289 * **ani** -- Animation.

290 '''

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

292 error_handling.positive_real(s, 's')

293 error_handling.positive_real(fs, 'fs')

294 error_handling.prop_type(prop_type)

295 error_handling.tuple_2elem(figsize, 'figsize')

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

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

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

299 ticks = [-max_val, max_val]

300 cmap = 'seismic'

301 else:

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

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

304 cmap = 'Reds'

305 fig = plt.figure(figsize=figsize)

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

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

308 ax.set_aspect('equal')

309 frame = plt.gca()

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

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

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

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

314 cmap=cmap)

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

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

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

318 else:

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

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

321

322 def init():

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

324 return scat,

325

326 def animate(i):

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

328 return scat,

329

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

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

332

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

334 '''

335 Plot time evolution for dimers.

336

337 :param lw: Positive number. Default value 5. Linewidth.

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

339

340 :returns:

341 * **fig** -- Figure.

342 '''

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

344 error_handling.positive_real(lw, 'lw')

345 error_handling.positive_real(fs, 'fs')

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

347 fig, ax = plt.subplots()

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

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

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

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

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

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

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

355 return fig

356

357

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

359propagation = Propagation