Coverage for tbkit/floquet.py: 100%

1r""" 

2Floquet theory of periodically driven Tight-Binding models, 

3:math:`H(t + T) = H(t)`, :math:`\omega = 2\pi/T` (:math:`\hbar = 1`). 

4 

5The evolution over one period, :math:`U(T) = \mathcal{T}e^{-i\int_0^TH(t)dt} 

6= e^{-iH_FT}`, defines the Floquet Hamiltonian :math:`H_F`, whose 

7eigenvalues -- the quasienergies :math:`\epsilon`, defined modulo 

8:math:`\omega` -- replace the energies of a static model: the driven 

9system, observed once per period, behaves as if it were static with 

10Hamiltonian :math:`H_F`. Driving can therefore *engineer* band structures 

11and their topology: circularly polarized light turns graphene into a 

12Chern insulator (Oka and Aoki, Phys. Rev. B 79, 081406 (2009); Kitagawa, 

13Oka, Brataas, Fu and Demler, Phys. Rev. B 84, 235108 (2011)). 

14 

15Two routes give the quasienergies, checked against each other in the tests: 

16 

17* the time-ordered product of short-time propagators (*evolution_operator*, 

18 *quasienergies*, *effective_hamiltonian*); 

19* the Sambe (extended) space of the harmonics 

20 :math:`H(t) = \sum_m H_m e^{-im\omega t}`, a static eigenvalue problem 

21 (*harmonics*, *sambe_hamiltonian*). 

22 

23:class:`FloquetKSpace` makes a driven **KSpace** model: its *get_ham* is 

24the Floquet Hamiltonian :math:`H_F(\mathbf{k})`, so its bands, Berry 

25curvature and Chern numbers are the driven ones. 

26 

27Anomalous Floquet phases. The quasienergies are defined modulo 

28:math:`\omega`, so :math:`H_F` depends on where the branch cut of the 

29logarithm is put (the *epsilon* arguments). The Chern numbers of the 

30Floquet bands no longer count the edge states: every band can have 

31:math:`C = 0` while chiral edge states cross every gap (Kitagawa, Berg, 

32Rudner and Demler, Phys. Rev. B 82, 235114 (2010)). The winding number of 

33Rudner, Lindner, Berg and Levin (Phys. Rev. X 3, 031005 (2013)), which 

34uses the evolution at *all* times of the period, does 

35(*DrivenKSpace.winding_number*): 

36 

37* :class:`DrivenKSpace` makes a driven Bloch model out of any 

38 time-dependent :math:`H(\mathbf{k}, t)`; 

39* *step_drive* builds a piecewise-constant drive from a list of models 

40 (**KSpace** models, ribbons made by *kspace.ribbon*, real-space 

41 **System** instances, or matrices), evolved exactly as a product of 

42 step exponentials (:class:`StepDrive` in real space); 

43* *DrivenKSpace.edge_state_count* counts the chiral edge states of a 

44 driven ribbon, independently of the bulk winding number. 

45""" 

46from __future__ import annotations 

47 

48import copy 

49from typing import TYPE_CHECKING, Callable 

50 

51import numpy as np 

52from numpy.typing import ArrayLike, NDArray 

53import scipy.linalg as LA 

54 

55import tbkit.error_handling as error_handling 

56 

57if TYPE_CHECKING: 

58 from tbkit.kspace import KSpace 

59 from tbkit.lattice import Lattice 

60 

61 

62def evolution_operator(ham_t: Callable, period: float, n_steps: int = 100) -> NDArray[np.complex128]: 

63 r''' 

64 Get the one-period evolution operator 

65 :math:`U(T) = \prod_m e^{-iH(t_m)\,\delta t}` (later times on the left, 

66 midpoints :math:`t_m = (m + 1/2)\delta t`, error :math:`O(\delta t^2)`). 

67 

68 Parameters 

69 ---------- 

70 ham_t : Callable 

71 ``ham_t(t)`` returns the Hamiltonian at time t (a square matrix). 

72 period : Positive real 

73 :math:`T`. 

74 n_steps : Positive integer, default 100 

75 Time steps per period. 

76 

77 Returns 

78 ------- 

79 U : Complex ndarray, the shape of the Hamiltonian 

80 ''' 

81 error_handling.is_callable(ham_t, 'ham_t') 

82 error_handling.positive_real(period, 'period') 

83 error_handling.positive_int(n_steps, 'n_steps') 

84 dt = period / n_steps 

85 u = LA.expm(-1j * dt * np.asarray(ham_t(0.5 * dt), dtype='c16')) 

86 for m in range(1, n_steps): 

87 h = np.asarray(ham_t((m + 0.5) * dt), dtype='c16') 

88 u = LA.expm(-1j * dt * h) @ u 

89 return u 

90 

91 

92def _fold(eps: NDArray[np.float64], period: float, epsilon: float) -> NDArray[np.float64]: 

93 ''' 

94 Private function. Quasienergies folded into [epsilon - 2 pi/T, epsilon). 

95 ''' 

96 omega = 2 * np.pi / period 

97 rest = np.mod(eps - epsilon, omega) 

98 rest = np.where(rest >= omega, rest - omega, rest) # np.mod may round up to omega 

99 return epsilon - omega + rest 

100 

101 

102def _floquet( 

103 u: NDArray[np.complex128], period: float, epsilon: float | None = None, 

104) -> tuple[NDArray[np.float64], NDArray[np.complex128]]: 

105 ''' 

106 Private function. Quasienergies (sorted, in [-pi/T, pi/T), or in 

107 [epsilon - 2 pi/T, epsilon)) and Floquet states of a unitary 

108 U = exp(-i H_F T), from its complex Schur form. 

109 ''' 

110 tri, z = LA.schur(u, output='complex') 

111 eps = -np.angle(np.diag(tri)) / period 

112 if epsilon is not None: 

113 eps = _fold(eps, period, epsilon) 

114 order = np.argsort(eps, kind='stable') 

115 return eps[order], z[:, order] 

116 

117 

118def _heff(u: NDArray[np.complex128], period: float, epsilon: float | None = None) -> NDArray[np.complex128]: 

119 ''' 

120 Private function. The Floquet Hamiltonian (i/T) ln U, branch cut at epsilon. 

121 ''' 

122 eps, z = _floquet(u, period, epsilon) 

123 h_f = z @ np.diag(eps) @ z.conj().T 

124 return (h_f + h_f.conj().T) / 2 

125 

126 

127def quasienergies( 

128 ham_t: Callable, period: float, n_steps: int = 100, epsilon: float | None = None, 

129) -> NDArray[np.float64]: 

130 r''' 

131 Get the quasienergies :math:`\epsilon`, eigenvalues of the Floquet 

132 Hamiltonian, in the first Floquet zone :math:`[-\omega/2, \omega/2)`, 

133 or in :math:`[\epsilon_{cut} - \omega, \epsilon_{cut})`. 

134 

135 Parameters 

136 ---------- 

137 ham_t : Callable 

138 See *evolution_operator*. 

139 period : Positive real 

140 n_steps : Positive integer, default 100 

141 epsilon : Real number 

142 Default value None (the zone 

143 :math:`[-\omega/2, \omega/2)`). Branch cut, see *effective_hamiltonian*. 

144 

145 Returns 

146 ------- 

147 eps : Real ndarray, sorted 

148 ''' 

149 error_handling.is_callable(ham_t, 'ham_t') 

150 error_handling.positive_real(period, 'period') 

151 error_handling.positive_int(n_steps, 'n_steps') 

152 error_handling.branch_cut(epsilon) 

153 return _floquet(evolution_operator(ham_t, period, n_steps), period, epsilon)[0] 

154 

155 

156def effective_hamiltonian( 

157 ham_t: Callable, period: float, n_steps: int = 100, epsilon: float | None = None, 

158) -> NDArray[np.complex128]: 

159 r''' 

160 Get the Floquet Hamiltonian :math:`H_F = \frac{i}{T}\ln U(T)`, with its 

161 quasienergies in the first Floquet zone (Hermitian for a Hermitian drive). 

162 

163 The logarithm is defined up to multiples of :math:`\omega` on each 

164 eigenvalue: *epsilon* puts its branch cut at the quasienergy 

165 :math:`\epsilon_{cut}` (the phase :math:`e^{-i\epsilon_{cut}T}` of 

166 :math:`U`), so that the quasienergies lie in 

167 :math:`[\epsilon_{cut} - \omega, \epsilon_{cut})`. The default, the zone 

168 :math:`[-\omega/2, \omega/2)`, is the cut :math:`\epsilon_{cut} = \omega/2`. 

169 Rudner et al. (2013) take :math:`(\epsilon_{cut}, \epsilon_{cut} + \omega]` 

170 instead: :math:`H_F` shifted by :math:`\omega` (the same projectors, 

171 so the same Chern numbers and winding numbers). 

172 

173 Parameters 

174 ---------- 

175 ham_t : Callable 

176 See *evolution_operator*. 

177 period : Positive real 

178 n_steps : Positive integer, default 100 

179 epsilon : Real number 

180 Default value None (:math:`\epsilon_{cut} = \omega/2`, 

181 computed exactly as before this option existed). Branch cut, 

182 best placed in a quasienergy gap. 

183 

184 Returns 

185 ------- 

186 H_F : Complex ndarray 

187 ''' 

188 error_handling.is_callable(ham_t, 'ham_t') 

189 error_handling.positive_real(period, 'period') 

190 error_handling.positive_int(n_steps, 'n_steps') 

191 error_handling.branch_cut(epsilon) 

192 return _heff(evolution_operator(ham_t, period, n_steps), period, epsilon) 

193 

194 

195def harmonics(ham_t: Callable, period: float, n_harmonics: int, n_samples: int | None = None) -> dict: 

196 r''' 

197 Get the Fourier components :math:`H_m = \frac1T\int_0^T H(t)\,e^{im\omega t}\,dt` 

198 of the drive, :math:`H(t) = \sum_m H_m e^{-im\omega t}`, for 

199 :math:`|m| \le` *n_harmonics* (exact for a drive with no higher harmonics). 

200 

201 Parameters 

202 ---------- 

203 ham_t : Callable 

204 See *evolution_operator*. 

205 period : Positive real 

206 n_harmonics : Positive integer or zero 

207 n_samples : Positive integer 

208 Default value None (4 n_harmonics + 16). 

209 Time samples per period. 

210 

211 Returns 

212 ------- 

213 H_m : Dictionary {m: complex ndarray} 

214 ''' 

215 error_handling.is_callable(ham_t, 'ham_t') 

216 error_handling.positive_real(period, 'period') 

217 error_handling.positive_int_zero(n_harmonics, 'n_harmonics') 

218 if n_samples is None: 

219 n_samples = 4 * n_harmonics + 16 

220 error_handling.positive_int(n_samples, 'n_samples') 

221 omega = 2 * np.pi / period 

222 ts = period * np.arange(n_samples) / n_samples 

223 hams = np.array([np.asarray(ham_t(t), dtype='c16') for t in ts]) 

224 return {m: np.mean(hams * np.exp(1j * m * omega * ts)[:, None, None], axis=0) 

225 for m in range(-n_harmonics, n_harmonics + 1)} 

226 

227 

228def sambe_hamiltonian(harm: dict, omega: float, n_max: int) -> NDArray[np.complex128]: 

229 r''' 

230 Get the Sambe-space Floquet Hamiltonian, truncated to the photon 

231 sectors :math:`n = -n_{max}, \dots, n_{max}`: 

232 

233 .. math:: 

234 

235 (\mathcal{H})_{nn'} = H_{n-n'} + n\omega\,\delta_{nn'}\, . 

236 

237 Its eigenvalues are the quasienergies, repeated in every Floquet zone 

238 (shifted by multiples of :math:`\omega`); those of the central zone 

239 converge fastest as *n_max* grows. 

240 

241 Parameters 

242 ---------- 

243 harm : Dictionary {m: square complex array}, see *harmonics* 

244 omega : Positive real 

245 Driving frequency. 

246 n_max : Positive integer or zero 

247 

248 Returns 

249 ------- 

250 H : Complex ndarray, shape ((2 n_max + 1) d, (2 n_max + 1) d) 

251 ''' 

252 error_handling.harmonics(harm) 

253 error_handling.positive_real(omega, 'omega') 

254 error_handling.positive_int_zero(n_max, 'n_max') 

255 d = len(harm[0]) 

256 size = 2 * n_max + 1 

257 out = np.zeros((size * d, size * d), 'c16') 

258 for a, n in enumerate(range(-n_max, n_max + 1)): 

259 for b, n2 in enumerate(range(-n_max, n_max + 1)): 

260 block = np.asarray(harm.get(n - n2, np.zeros((d, d))), dtype='c16') 

261 if n == n2: 

262 block = block + n * omega * np.eye(d) 

263 out[a*d:(a+1)*d, b*d:(b+1)*d] = block 

264 return out 

265 

266 

267# ---------------------------------------------------------------------- 

268# Piecewise-constant (step) drives 

269# ---------------------------------------------------------------------- 

270 

271 

272def _step_evolution(steps: list, t: float | None = None) -> NDArray[np.complex128]: 

273 r''' 

274 Private function. The exact evolution of a piecewise-constant drive, 

275 *steps* = [(H_1, dt_1), (H_2, dt_2), ...] in time order: 

276 U(t) = exp(-i H_n (t - t_{n-1})) ... exp(-i H_1 dt_1), U(T) if t is None. 

277 ''' 

278 u = np.eye(len(steps[0][0]), dtype='c16') 

279 elapsed = 0. 

280 for h, dt in steps: 

281 tau = dt if t is None else min(dt, t - elapsed) 

282 if tau <= 0: 

283 break 

284 u = LA.expm(-1j * tau * h) @ u 

285 elapsed += dt 

286 return u 

287 

288 

289class StepDrive: 

290 r''' 

291 A finite (real-space) system driven by a piecewise-constant 

292 Hamiltonian: :math:`H(t) = H_n` for 

293 :math:`t_{n-1} \le t < t_n = \delta t_1 + \dots + \delta t_n`, repeated 

294 with period :math:`T = \sum_n \delta t_n`. The evolution is exact, a 

295 product of step exponentials, 

296 

297 .. math:: 

298 

299 U(t) = e^{-iH_n(t - t_{n-1})}\,e^{-iH_{n-1}\delta t_{n-1}}\cdots e^{-iH_1\delta t_1}\, , 

300 

301 with no time grid. Built by *step_drive* from **System** instances or 

302 matrices (see *DrivenKSpace* for Bloch models and ribbons). 

303 

304 Parameters 

305 ---------- 

306 hams 

307 List of square complex ndarrays, one Hermitian matrix per step (all of one shape). 

308 durations : List of positive reals, the step durations :math:`\delta t_n` 

309 epsilon : Real number, default None 

310 Default branch cut of 

311 *quasienergies* and *effective_hamiltonian*. 

312 ''' 

313 

314 def __init__(self, hams: list, durations: ArrayLike, epsilon: float | None = None) -> None: 

315 error_handling.step_matrices(hams) 

316 error_handling.durations(durations, len(hams)) 

317 error_handling.branch_cut(epsilon) 

318 self.hams = [np.asarray(h, dtype='c16') for h in hams] 

319 self.durations = np.array(durations, dtype='f8') 

320 self.period = float(self.durations.sum()) 

321 self.epsilon = epsilon 

322 self.n_sites = len(self.hams[0]) 

323 

324 def get_ham(self, t: float) -> NDArray[np.complex128]: 

325 r''' 

326 Get the Hamiltonian :math:`H(t)` (periodic in :math:`T`). 

327 

328 Parameters 

329 ---------- 

330 t : Real number 

331 Time. 

332 

333 Returns 

334 ------- 

335 H : Complex ndarray 

336 ''' 

337 error_handling.real_number(t, 't') 

338 n = int(np.searchsorted(np.cumsum(self.durations), t % self.period, side='right')) 

339 return self.hams[min(n, len(self.hams) - 1)] 

340 

341 def evolution_operator(self, t: float | None = None) -> NDArray[np.complex128]: 

342 r''' 

343 Get the exact evolution operator :math:`U(t)` from time 0. 

344 

345 Parameters 

346 ---------- 

347 t : Real number in :math:`[0, T]` 

348 Default value None (one period, :math:`U(T)`). 

349 

350 Returns 

351 ------- 

352 U : Complex ndarray 

353 ''' 

354 error_handling.time_in_period(t, self.period) 

355 return _step_evolution(list(zip(self.hams, self.durations)), t) 

356 

357 def quasienergies(self, epsilon: float | None = None) -> NDArray[np.float64]: 

358 r''' 

359 Get the quasienergies, sorted, in :math:`[-\omega/2, \omega/2)` or 

360 :math:`[\epsilon_{cut} - \omega, \epsilon_{cut})` (see *effective_hamiltonian*). 

361 

362 Parameters 

363 ---------- 

364 epsilon : Real number 

365 Default value None (the *epsilon* of the drive). 

366 

367 Returns 

368 ------- 

369 eps : Real ndarray 

370 ''' 

371 error_handling.branch_cut(epsilon) 

372 epsilon = self.epsilon if epsilon is None else epsilon 

373 return _floquet(self.evolution_operator(), self.period, epsilon)[0] 

374 

375 def effective_hamiltonian(self, epsilon: float | None = None) -> NDArray[np.complex128]: 

376 r''' 

377 Get the Floquet Hamiltonian :math:`H_F = \frac{i}{T}\ln U(T)`, branch 

378 cut at *epsilon* (see *tbkit.floquet.effective_hamiltonian*). 

379 

380 Parameters 

381 ---------- 

382 epsilon : Real number 

383 Default value None (the *epsilon* of the drive). 

384 

385 Returns 

386 ------- 

387 H_F : Complex ndarray 

388 ''' 

389 error_handling.branch_cut(epsilon) 

390 epsilon = self.epsilon if epsilon is None else epsilon 

391 return _heff(self.evolution_operator(), self.period, epsilon) 

392 

393 def evolve(self, psi: NDArray, n_periods: int) -> NDArray[np.complex128]: 

394 r''' 

395 Evolve a state stroboscopically, :math:`\psi(nT) = U(T)^n\psi(0)`. 

396 

397 Parameters 

398 ---------- 

399 psi : Complex ndarray, the initial state (one amplitude per site) 

400 n_periods : Positive integer or zero 

401 

402 Returns 

403 ------- 

404 psi_n 

405 Complex ndarray, shape (n_periods + 1, n_sites): 

406 the state after 0, 1, ..., n_periods periods. 

407 ''' 

408 error_handling.ndarray(psi, 'psi', self.n_sites) 

409 error_handling.positive_int_zero(n_periods, 'n_periods') 

410 u = self.evolution_operator() 

411 out = np.zeros((n_periods + 1, self.n_sites), 'c16') 

412 out[0] = psi 

413 for n in range(n_periods): 

414 out[n + 1] = u @ out[n] 

415 return out 

416 

417 

418def step_drive(models: list, durations: ArrayLike, epsilon: float | None = None): 

419 r''' 

420 Build a piecewise-constant drive: the Hamiltonian of ``models[n]`` acts 

421 for ``durations[n]``, in order, and the cycle repeats with period 

422 :math:`T = \sum_n \delta t_n`. The evolution is exact (a product of step 

423 exponentials, no time grid, see *StepDrive*). 

424 

425 Parameters 

426 ---------- 

427 models 

428 List of models of one size, all of one kind: 

429 

430 * **KSpace** instances (including ribbons made by *kspace.ribbon*) 

431 -- returns a *DrivenKSpace*, whose *get_ham* is the Floquet 

432 Hamiltonian :math:`H_F(\mathbf{k})`; 

433 * **System** instances (their *get_ham* is called) or square 

434 matrices (dense or scipy.sparse) -- returns a *StepDrive*. 

435 

436 Each must be Hermitian; KSpace models need an orthogonal basis. 

437 durations : List of positive reals 

438 :math:`\delta t_n`, one per model. 

439 epsilon : Real number, default None 

440 Branch cut of the 

441 Floquet Hamiltonian (see *effective_hamiltonian*). 

442 

443 Returns 

444 ------- 

445 drive 

446 *DrivenKSpace* or *StepDrive*. 

447 

448 Examples 

449 -------- 

450 Example usage:: 

451 

452 # Rudner et al. (2013): hopping along one bond direction per step 

453 drive = step_drive([ks1, ks2, ks3, ks4, ks5], [T/5] * 5) 

454 w = drive.winding_number(epsilon=np.pi / T) 

455 ''' 

456 from tbkit.kspace import KSpace 

457 from tbkit.system import System 

458 kind = error_handling.step_models(models, KSpace, System) 

459 error_handling.durations(durations, len(models)) 

460 error_handling.branch_cut(epsilon) 

461 if kind == 'kspace': 

462 return DrivenKSpace._from_steps(models, durations, epsilon) 

463 hams = [] 

464 for m in models: 

465 if isinstance(m, System): 

466 m.get_ham() 

467 m = m.ham 

468 hams.append(m.toarray() if hasattr(m, 'toarray') else np.asarray(m)) 

469 return StepDrive(hams, durations, epsilon) 

470 

471 

472# ---------------------------------------------------------------------- 

473# Driven Bloch models 

474# ---------------------------------------------------------------------- 

475 

476 

477def _bloch_grid(model, ks: NDArray[np.float64], a: ArrayLike | None = None) -> NDArray[np.complex128]: 

478 r''' 

479 Private function. The Bloch Hamiltonians of a static **KSpace** at many 

480 k-points at once, shape (nk, norb, norb): *get_ham* (or, with a vector 

481 potential *a*, *get_ham_peierls*) vectorized over k. A model whose class 

482 redefines these methods is evaluated point by point. 

483 ''' 

484 from tbkit.kspace import KSpace 

485 if type(model).get_ham is not KSpace.get_ham or type(model).get_ham_peierls is not KSpace.get_ham_peierls: 

486 if a is None: 

487 return np.array([model.get_ham(k) for k in ks]) 

488 return np.array([model.get_ham_peierls(k, a) for k in ks]) 

489 return model._bloch_ham(ks @ model.k_basis.T, None if a is None else np.asarray(a, dtype='f8')) 

490 

491 

492def _spectral_derivative(u: NDArray[np.complex128], axis: int) -> NDArray[np.complex128]: 

493 r''' 

494 Private function. Derivative of a function sampled on a uniform grid 

495 of one period [0, 1) along *axis* (exact for a trigonometric 

496 polynomial, exponentially accurate for an analytic periodic function). 

497 ''' 

498 n = u.shape[axis] 

499 m = np.fft.fftfreq(n, 1. / n) 

500 if n % 2 == 0: 

501 m[n // 2] = 0. # the Nyquist mode has no odd derivative 

502 shape = [1] * u.ndim 

503 shape[axis] = n 

504 return np.fft.ifft(np.fft.fft(u, axis=axis) * (2j * np.pi * m).reshape(shape), axis=axis) 

505 

506 

507def _winding_density(u: NDArray[np.complex128], a_t: NDArray[np.complexfloating]) -> float: 

508 r''' 

509 Private function. Brillouin-zone average of 

510 Tr(A_t [A_1, A_2]), A_i = U^-1 d U / d f_i along the fractional 

511 coordinates f_1, f_2 (u has shape (n1, n2, N, N)). 

512 ''' 

513 u_inv = np.conj(np.swapaxes(u, -1, -2)) 

514 a1 = u_inv @ _spectral_derivative(u, 0) 

515 a2 = u_inv @ _spectral_derivative(u, 1) 

516 return float(np.mean(np.trace(a_t @ (a1 @ a2 - a2 @ a1), axis1=-2, axis2=-1)).real) 

517 

518 

519class DrivenKSpace: 

520 r''' 

521 A periodic model driven by any time-periodic Bloch Hamiltonian 

522 :math:`H(\mathbf{k}, t + T) = H(\mathbf{k}, t)`. 

523 

524 It is a copy of the **KSpace** model whose *get_ham* returns the Floquet 

525 Hamiltonian :math:`H_F(\mathbf{k}) = \frac{i}{T}\ln U(\mathbf{k}, T)` 

526 (branch cut at *epsilon*, see *effective_hamiltonian*), so *get_bands*, 

527 *k_path*, *berry_curvature*, *chern_number*, ... describe the driven 

528 system, and *winding_number* gives its anomalous topology. The 

529 evolution :math:`U(\mathbf{k}, T)` is the product of *n_steps* 

530 midpoint propagators (see *evolution_operator*); drives built by 

531 *step_drive* are evolved exactly. Real-space methods that need static 

532 hoppings (*finite_ham*, *get_ham_beta*, *gbz*) are not available. 

533 

534 Parameters 

535 ---------- 

536 ks : **KSpace** instance 

537 The model driven: its lattice, orbitals 

538 and k-space conventions (periodic gauge) are those of the driven model. 

539 ham_kt : Callable 

540 ``ham_kt(k, t)`` returns the Hermitian Bloch 

541 Hamiltonian :math:`H(\mathbf{k}, t)`, shape (norb, norb), in the 

542 periodic gauge of *ks* (:math:`H(\mathbf{k} + \mathbf{G}, t) = H(\mathbf{k}, t)`). 

543 period : Positive real 

544 :math:`T`. 

545 n_steps : Positive integer, default 100 

546 Time steps per period. 

547 epsilon : Real number 

548 Default value None (quasienergies in 

549 :math:`[-\omega/2, \omega/2)`). Branch cut of *get_ham*. 

550 

551 Examples 

552 -------- 

553 Example usage:: 

554 

555 # a sublattice mass switched on and off 

556 drive = DrivenKSpace(ks, lambda k, t: ks.get_ham(k) + m * np.cos(w * t) * sz, 2 * np.pi / w) 

557 ''' 

558 

559 # Set by _copy (and _from_steps or __new__) on an instance whose class 

560 # also derives from type(ks): the rest is the copied KSpace's. 

561 static: KSpace 

562 period: float 

563 epsilon: float | None 

564 steps: list | None 

565 durations: NDArray[np.float64] 

566 n_steps: int 

567 ham_kt: Callable 

568 lat: Lattice 

569 dim: int 

570 norb: int 

571 spin: bool 

572 rec_vec_k: NDArray[np.float64] 

573 _plane_orientation: Callable 

574 

575 def __new__(cls, ks, ham_kt: Callable, period: float, n_steps: int = 100, epsilon: float | None = None): 

576 from tbkit.kspace import KSpace 

577 error_handling.kspace(ks, KSpace) 

578 error_handling.is_callable(ham_kt, 'ham_kt') 

579 error_handling.positive_real(period, 'period') 

580 error_handling.positive_int(n_steps, 'n_steps') 

581 error_handling.branch_cut(epsilon) 

582 error_handling.drive_matrix(ham_kt(np.zeros(ks.dim), 0.), ks.norb) 

583 obj = cls._copy(ks, period, epsilon) 

584 obj.ham_kt = ham_kt 

585 obj.n_steps = n_steps 

586 return obj 

587 

588 def __init__(self, ks, ham_kt, period, n_steps=100, epsilon=None) -> None: 

589 pass 

590 

591 @classmethod 

592 def _copy(cls, ks, period: float, epsilon: float | None): 

593 ''' 

594 Private method. A copy of *ks* whose class also derives from *cls*. 

595 ''' 

596 driven_cls = type(cls.__name__, (cls, type(ks)), {}) 

597 obj: DrivenKSpace = object.__new__(driven_cls) 

598 obj.__dict__.update(copy.deepcopy(ks.__dict__)) 

599 obj.static = ks 

600 obj.period = period 

601 obj.epsilon = epsilon 

602 obj.steps = None 

603 return obj 

604 

605 @classmethod 

606 def _from_steps(cls, models: list, durations: ArrayLike, epsilon: float | None): 

607 ''' 

608 Private method. The drive of *step_drive*: models[n] for durations[n]. 

609 ''' 

610 durations = np.array(durations, dtype='f8') 

611 obj = cls._copy(models[0], float(durations.sum()), epsilon) 

612 obj.steps = list(models) 

613 obj.durations = durations 

614 obj.n_steps = len(models) 

615 obj.ham_kt = obj._step_ham 

616 return obj 

617 

618 def _step_ham(self, k: ArrayLike, t: float) -> NDArray[np.complex128]: 

619 ''' 

620 Private method. H(k, t) of a step drive. 

621 ''' 

622 steps = self.steps or [] # set by _from_steps, the only caller 

623 n = int(np.searchsorted(np.cumsum(self.durations), t % self.period, side='right')) 

624 return steps[min(n, len(steps) - 1)].get_ham(k) 

625 

626 def _step_list(self, k: NDArray[np.float64]) -> list: 

627 ''' 

628 Private method. The drive at k as steps [(H_n(k), dt_n)]: exact for 

629 a step drive, the midpoint propagators otherwise. 

630 ''' 

631 if self.steps is not None: 

632 return [(m.get_ham(k), d) for m, d in zip(self.steps, self.durations)] 

633 dt = self.period / self.n_steps 

634 return [(np.asarray(self.ham_kt(k, (m + 0.5) * dt), dtype='c16'), dt) for m in range(self.n_steps)] 

635 

636 def _grid_steps(self, ks: NDArray[np.float64]) -> tuple[NDArray[np.complex128], list]: 

637 ''' 

638 Private method. The steps of *_step_list* at many k-points at once: 

639 the Hamiltonians, shape (nk, n_steps, norb, norb), and the durations. 

640 ''' 

641 if self.steps is not None: 

642 return np.stack([_bloch_grid(m, ks) for m in self.steps], axis=1), list(self.durations) 

643 dt = self.period / self.n_steps 

644 hams = np.array([[s[0] for s in self._step_list(k)] for k in ks]) 

645 return hams, [dt] * self.n_steps 

646 

647 def get_evolution(self, k: ArrayLike, t: float | None = None) -> NDArray[np.complex128]: 

648 r''' 

649 Get the evolution operator :math:`U(\mathbf{k}, t) = \mathcal{T}e^{-i\int_0^tH(\mathbf{k},t')dt'}`. 

650 

651 Parameters 

652 ---------- 

653 k : k point (see *KSpace.get_ham*) 

654 t : Real number in :math:`[0, T]` 

655 Default value None (one 

656 period: the Floquet operator :math:`U(\mathbf{k}, T)`). 

657 

658 Returns 

659 ------- 

660 U : Complex ndarray, shape (norb, norb) 

661 ''' 

662 error_handling.k_vector(k, 'k', self.dim) 

663 error_handling.time_in_period(t, self.period) 

664 k = np.asarray(k, dtype='f8') 

665 if t is None and self.steps is None: 

666 return evolution_operator(lambda tt: self.ham_kt(k, tt), self.period, self.n_steps) 

667 return _step_evolution(self._step_list(k), t) 

668 

669 def get_ham(self, k: ArrayLike) -> NDArray[np.complex128]: 

670 r''' 

671 Get the Floquet Hamiltonian :math:`H_F(\mathbf{k})`, quasienergies in 

672 :math:`[\epsilon_{cut} - \omega, \epsilon_{cut})` (default 

673 :math:`[-\omega/2, \omega/2)`). 

674 

675 Parameters 

676 ---------- 

677 k : k point (see *KSpace.get_ham*) 

678 

679 Returns 

680 ------- 

681 H_F : Complex ndarray, shape (norb, norb) 

682 ''' 

683 error_handling.k_vector(k, 'k', self.dim) 

684 return self._hams(np.asarray(k, dtype='f8')[None])[0] 

685 

686 def _hams(self, ks: NDArray[np.float64]) -> NDArray[np.complex128]: 

687 ''' 

688 Private method. *get_ham* at the k-points *ks* (shape (nk, dim)): the 

689 evolutions of all the k-points advance together, one stacked 

690 exponential per time step (the steps of *_grid_steps*), in chunks of 

691 bounded memory. 

692 ''' 

693 chunk = max(1, 400000 // (self.norb ** 2 * self.n_steps)) 

694 out = [] 

695 for c0 in range(0, len(ks), chunk): 

696 hams, durations = self._grid_steps(ks[c0:c0 + chunk]) 

697 u = np.broadcast_to(np.eye(self.norb, dtype='c16'), (len(hams), self.norb, self.norb)) 

698 for s, dt in enumerate(durations): 

699 u = LA.expm(-1j * dt * hams[:, s]) @ u 

700 out += [_heff(uk, self.period, self.epsilon) for uk in u] 

701 return np.array(out, dtype='c16').reshape(len(ks), self.norb, self.norb) 

702 

703 def is_hermitian(self) -> bool: 

704 ''' 

705 The Floquet Hamiltonian of a Hermitian drive is Hermitian. 

706 ''' 

707 return True 

708 

709 def _hop_cells(self): 

710 error_handling.not_floquet() 

711 

712 def _check_static(self): 

713 error_handling.not_floquet() 

714 

715 def winding_number(self, epsilon: float, nk: int | tuple[int, int] = 48, n_t: int = 40) -> float: 

716 r''' 

717 Get the winding number of Rudner, Lindner, Berg and Levin (Phys. Rev. 

718 X 3, 031005 (2013)), the number of chiral edge states in the 

719 quasienergy gap at *epsilon*: 

720 

721 .. math:: 

722 

723 W[U_\epsilon] = \frac{1}{8\pi^2}\int_0^T dt\int_{BZ} dk_x dk_y\, 

724 \mathrm{Tr}\left(U_\epsilon^{-1}\partial_tU_\epsilon 

725 \left[U_\epsilon^{-1}\partial_{k_x}U_\epsilon, U_\epsilon^{-1}\partial_{k_y}U_\epsilon\right]\right), 

726 

727 of the loop :math:`U_\epsilon(\mathbf{k}, t) = U(\mathbf{k}, 2t)` for 

728 :math:`t \le T/2`, then :math:`V_\epsilon(\mathbf{k}, 2T - 2t)`, with 

729 the return map :math:`V_\epsilon(\mathbf{k}, t) = e^{-iH_F^\epsilon(\mathbf{k})t}` 

730 and :math:`H_F^\epsilon` the Floquet Hamiltonian with its branch cut 

731 at *epsilon* (see *effective_hamiltonian*). As the integrand is a 

732 3-form, this is the winding of :math:`U` over :math:`[0, T]` minus 

733 that of :math:`V_\epsilon` over :math:`[0, T]`. 

734 

735 Signs. :math:`W` counts the chiral edge states in the gap: 

736 :math:`W = n_{lower} = -n_{upper}` of *edge_state_count*, so 

737 :math:`W > 0` means states moving along :math:`+\mathbf{a}_1` on 

738 the edge of a ribbon on the side of :math:`-\hat{z}\times\mathbf{a}_1` 

739 (the bottom edge of a ribbon along x). Across a group of Floquet bands, 

740 :math:`W[U_{\epsilon'}] - W[U_\epsilon] = -\sum_n C_n` for 

741 :math:`\epsilon < \epsilon'`, the sum over the bands between the two 

742 gaps, with the Chern numbers :math:`C_n` of *chern_number* (the 

743 paper's :math:`\mathcal{C}_n = -\frac{1}{2\pi}\int\nabla\times\mathcal{A}_n` 

744 is minus ours, so there :math:`W' - W = +\sum\mathcal{C}`). A 

745 static model (whose bands fit in the zone below the cut) has 

746 :math:`W = 0` at the zone edge, and :math:`W = -\sum C` of the 

747 bands below any other gap. The anomalous phases are those with 

748 :math:`W \ne 0` in every gap while all the :math:`C_n = 0`. 

749 

750 Discretization: :math:`U(\mathbf{k}, t)` on a uniform 

751 :math:`n_1\times n_2` Brillouin-zone grid, differentiated spectrally 

752 in :math:`\mathbf{k}` (exponentially accurate: both :math:`U` and, 

753 while the gap at *epsilon* is open, :math:`V_\epsilon` are analytic 

754 and periodic in :math:`\mathbf{k}`), and :math:`\partial_tU = -iHU` 

755 exactly, integrated by Gauss-Legendre quadrature on each time step 

756 (*n_t* nodes per period in total, at least 2 per step) and over 

757 :math:`[0, T]` for the return map. Step drives are exact; a 

758 continuous drive is replaced by its *n_steps* midpoint steps, whose 

759 loop has the same :math:`W` as long as the gap stays open. 

760 

761 Tolerance. The result is returned raw. Its error falls off 

762 exponentially with *nk*, at a rate set by the k-width of the 

763 region where :math:`V_\epsilon` varies fast, gap / band velocity 

764 near the smallest gap at *epsilon* (the Berry curvature is 

765 concentrated there). For the five-step model 

766 (:math:`\delta_{AB} = 0.5\pi/T`) the error is below 

767 :math:`10^{-12}` at perfect transfer; with a gap of 

768 :math:`0.067\omega` it is :math:`2\times10^{-5}` at the default 

769 *nk* = 48; with a gap of :math:`0.021\omega` it is 

770 :math:`10^{-2}` at *nk* = 48 and :math:`4\times10^{-4}` at 

771 *nk* = 96. Circularly driven graphene (gap :math:`0.01\omega` at K) 

772 needs *nk* = 96 for :math:`3\times10^{-3}`. Rounding to the nearest 

773 integer is trustworthy when :math:`|W - \mathrm{round}(W)| < 0.01` and the 

774 value is unchanged within that tolerance when *nk* is doubled; a 

775 value far from an integer means the gap is too small for the grid 

776 (or closed). 

777 

778 Parameters 

779 ---------- 

780 epsilon : Real number 

781 Quasienergy of the gap (branch cut). 

782 nk : Positive integer (at least 4), or tuple of 2 

783 Default 

784 value 48. k-points along each reciprocal vector. 

785 n_t : Positive integer, default 40 

786 Gauss-Legendre 

787 nodes per period (at least 2 per step; 2 n_t, and at least 24, 

788 for the return map). 

789 

790 Returns 

791 ------- 

792 W : Real number, close to an integer 

793 ''' 

794 error_handling.dim_exact(self.dim, 2) 

795 error_handling.real_number(epsilon, 'epsilon') 

796 error_handling.nk(nk, 2) 

797 if isinstance(nk, int): 

798 nk = (nk, nk) 

799 error_handling.nk_min(nk, 4) 

800 error_handling.positive_int(n_t, 'n_t') 

801 n1, n2 = nk 

802 fracs = np.stack(np.meshgrid(np.arange(n1) / n1, np.arange(n2) / n2, indexing='ij'), -1) 

803 ks = fracs @ self.rec_vec_k 

804 hams, durations = self._grid_steps(ks.reshape(-1, 2)) 

805 hams = hams.reshape(n1, n2, len(durations), self.norb, self.norb) 

806 x, w = np.polynomial.legendre.leggauss(max(2, -(-n_t // len(durations)))) 

807 # the driven evolution U(k, t), 0 <= t <= T 

808 u = np.broadcast_to(np.eye(self.norb, dtype='c16'), (n1, n2, self.norb, self.norb)) 

809 integral = 0. 

810 for s, dt in enumerate(durations): 

811 h = hams[:, :, s] 

812 e, q = np.linalg.eigh(h) 

813 qh = np.conj(np.swapaxes(q, -1, -2)) 

814 for xj, wj in zip(x, w): 

815 ut = (q * np.exp(-1j * e * dt * (1 + xj) / 2)[..., None, :]) @ qh @ u 

816 a_t = -1j * np.conj(np.swapaxes(ut, -1, -2)) @ h @ ut 

817 integral += wj * dt / 2 * _winding_density(ut, a_t) 

818 u = (q * np.exp(-1j * e * dt)[..., None, :]) @ qh @ u 

819 # the return map V(k, t) = exp(-i H_F t), 0 <= t <= T, run backwards 

820 eps = np.zeros((n1, n2, self.norb)) 

821 z = np.zeros((n1, n2, self.norb, self.norb), 'c16')