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 '''

23 Get lattice time evolution. Time dependent Schrodinger equation solved by

24 Crank-Nicolson method.

26 :param lat: **lattice** class instance.

27 '''

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

30 error_handling.lat(lat)

31 self.lat = lat

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

34 def get_propagation(

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

36 ) -> None:

37 '''

38 Get the time evolution.

40 :param ham: sparse.csr_matrix. Tight-Binding Hamilonian.

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

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

43 :param dz: Positive number. Step.

44 :param norm: Boolean. Default value True. Normalize the norm to 1 at each step.

45 '''

46 error_handling.empty_ham(ham)

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

48 error_handling.positive_int(steps, 'steps')

49 error_handling.positive_real(dz, 'dz')

50 error_handling.boolean(norm, 'norm')

51 self.steps = steps

52 self.dz = dz

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

54 self.prop[:, 0] = psi_init

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

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

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

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

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

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

61 if norm:

62 self.prop[:, i] /= np.abs(self.prop[:, i]).sum()

64 def get_pumping(

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

66 ) -> None:

67 '''

68 Get the time evolution with adiabatic pumpings.

70 :param hams: List of sparse.csr_matrices. Tight-Binding Hamilonians.

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

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

73 :param dz: Positive number. Step.

74 :param norm: Boolean. Default value True. Normalize the norm to 1 at each step.

75 '''

76 error_handling.get_pump(hams)

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

78 error_handling.positive_int(steps, 'steps')

79 error_handling.positive_real(dz, 'dz')

80 error_handling.boolean(norm, 'norm')

81 self.steps = steps

82 self.dz = dz

83 no = len(hams)

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

85 self.prop[:, 0] = psi_init

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

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

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

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

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

91 # before pumping

92 for i in range(1, delta):

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

94 if norm:

95 self.prop[:, i] /= np.abs(self.prop[:, i]).sum()

96 # pumping

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

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

99 for i in range(0, delta):

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

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

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

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

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

105 if norm:

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

107 np.abs(self.prop[:, (j+1)*delta+i]).sum()

108 # after pumping

109 j = no

110 for i in range(0, self.steps - no*delta):

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

112 if norm:

113 self.prop[:, no*delta+i] /= np.abs(self.prop[:, no*delta+i]).sum()

114

115 def plt_propagation_1d(

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

117 ) -> Figure:

118 '''

119 Plot time evolution for 1D systems.

120

121 :param fs: Default value 20. Fontsize.

122 '''

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

124 error_handling.positive_real(fs, 'fs')

125 error_handling.prop_type(prop_type)

126 error_handling.tuple_2elem(figsize, 'figsize')

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

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

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

130 if prop_type == 'real':

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

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

133 ticks = [-max_val, max_val]

134 cmap = 'seismic'

135 elif prop_type == 'imag':

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

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

138 ticks = [-max_val, max_val]

139 cmap = 'seismic'

140 else:

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

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

143 cmap = plt.cm.hot

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

145 aspect = 'auto'

146 interpolation = 'nearest'

147 im = plt.imshow(color, cmap=cmap, aspect=aspect,

148 interpolation=interpolation, extent=extent,

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

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

151 label.set_fontsize(fs)

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

153 label.set_fontsize(fs)

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

155 if prop_type == 'norm':

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

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

158 else:

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

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

161 cbar.ax.tick_params(labelsize=fs)

162 return fig

163

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

165 r'''

166 Private function. Used in *plt_propagation_1d*.

167 Smooth propagation for 1D systems.

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

169

170 :param prop: Propagation.

171 :param a: Default value 15. Gaussian Parameter.

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

173

174 :returns:

175 * **smooth** -- Smoothed propagation.

176 '''

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

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

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

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

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

182 return smooth

183

184 def get_animation(

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

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

187 ) -> FuncAnimation:

188 '''

189 Get time evolution animation.

190

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

192 :param fs: Default value 20. Fontsize.

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

194 :param prop_type: Default value None. Figsize.

195

196 :returns:

197 * **ani** -- Animation.

198 '''

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

200 error_handling.positive_real(s, 's')

201 error_handling.positive_real(fs, 'fs')

202 error_handling.prop_type(prop_type)

203 error_handling.tuple_2elem(figsize, 'figsize')

204 if os.name == 'posix':

205 blit = False

206 else:

207 blit = True

208 if prop_type == 'real':

209 color = self.prop.real

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

211 ticks = [-max_val, max_val]

212 cmap = 'seismic'

213 elif prop_type == 'imag':

214 color = self.prop.imag

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

216 ticks = [-max_val, max_val]

217 cmap = 'seismic'

218 else:

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

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

221 cmap = 'Reds'

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

223 plt.xlim([self.lat.coor['x'][0]-1., self.lat.coor['x'][-1]+1.])

224 plt.ylim([self.lat.coor['y'][0]-1., self.lat.coor['y'][-1]+1.])

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

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

227 cmap=plt.get_cmap(cmap))

228 frame = plt.gca()

229 frame.axes.get_xaxis().set_ticks([])

230 frame.axes.get_yaxis().set_ticks([])

231 ax.set_aspect('equal')

232 if prop_type == 'norm':

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

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

235 else:

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

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

238

239 def update(i, color, scat):

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

241 return scat,

242

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

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

245 return ani

246

247 def get_animation_nb(

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

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

250 ) -> FuncAnimation:

251 '''

252 Get time evolution animation for iPython notebooks.

253

254 :param s: Default value 300. Circle shape.

255 :param fs: Default value 20. Fontsize.

256

257 :returns:

258 * **ani** -- Animation.

259 '''

260 '''

261 Get time evolution animation.

262

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

264 :param fs: Default value 20. Fontsize.

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

266 :param prop_type: Default value None. Figsize.

267

268 :returns:

269 * **ani** -- Animation.

270 '''

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

272 error_handling.positive_real(s, 's')

273 error_handling.positive_real(fs, 'fs')

274 error_handling.prop_type(prop_type)

275 error_handling.tuple_2elem(figsize, 'figsize')

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

277 color = self.prop.real

278 max_val = max(np.max(color[:, -1]), -np.min(color[:, -1]))

279 ticks = [-max_val, max_val]

280 cmap = 'seismic'

281 else:

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

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

284 cmap = 'Reds'

285 fig = plt.figure()

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

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

288 ax.set_aspect('equal')

289 frame = plt.gca()

290 frame.axes.get_xaxis().set_ticks([])

291 frame.axes.get_yaxis().set_ticks([])

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=cmap)

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

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

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

298 else:

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

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

301

302 def init():

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

304 return scat,

305

306 def animate(i):

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

308 return scat,

309

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

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

312

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

314 '''

315 Plot time evolution for dimers.

316

317 :param lw: Default value 5. Linewidth.

318 :param fs: Default value 20. Fontsize.

319

320 :returns:

321 * **fig** -- Figure.

322 '''

323 if not self.prop.any():

324 raise Exception('\n\nRun method get_prop() or get_pump() first.\n')

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

326 fig, ax = plt.subplots()

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

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

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

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

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

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

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

334 return fig

335

336

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

338propagation = Propagation