Coverage for tbkit/orbital.py: 100%
243 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 11:04 +0000
1r"""
2Real-space Tight-Binding models with several orbitals per site and an
3optional spin: :class:`OrbitalSystem` extends :class:`tbkit.system.System`
4with Slater-Koster hoppings between s, p and d orbitals (see
5:mod:`tbkit.slater_koster`), non-orthogonal bases (an overlap matrix),
6and spin-dependent terms -- Kane-Mele intrinsic spin-orbit coupling,
7Rashba coupling, Zeeman fields and atomic :math:`\lambda\mathbf{L}\cdot\mathbf{S}`.
9Rows of the Hamiltonian are ordered site by site, then orbital by orbital
10(in the order given for the site's tag), then spin (up, down): the row of
11orbital *o* with spin *s* on site *i* is *sys.row(i, o, s)*.
12"""
13from __future__ import annotations
15from typing import Callable
17import numpy as np
18from numpy.typing import NDArray
19import scipy.linalg as LA
20import scipy.sparse as sparse
22import tbkit.error_handling as error_handling
23from tbkit.lattice import Lattice
24from tbkit.slater_koster import sk_block, orbital_angular_momentum
25from tbkit.system import System
27SIGMA = [np.array([[0., 1.], [1., 0.]], dtype='c16'),
28 np.array([[0., -1j], [1j, 0.]], dtype='c16'),
29 np.array([[1., 0.], [0., -1.]], dtype='c16')]
32class OrbitalSystem(System):
33 r'''
34 Solve a real-space Tight-Binding model with several orbitals per site
35 and, optionally, spin.
37 Parameters
38 ----------
39 lat : **Lattice** class instance
40 orbitals : Dictionary
41 Default value None (one 's' orbital per
42 site). key: sublattice tag, val: list of orbital names (see
43 :data:`tbkit.slater_koster.ORBITALS`), e.g. ``{'a': ['s', 'px', 'py', 'pz']}``.
44 spin : Boolean, default False
45 Give every orbital a
46 spin-1/2 degree of freedom.
47 periodic : Boolean, or tuple of booleans, default False
48 Periodic boundaries (see *System*); the bond vectors of the
49 Slater-Koster, spin-orbit and Rashba blocks are then those of the
50 shortest image.
52 Notes
53 -----
54 Everything of **System** still applies: *set_hopping* (and the disorder,
55 defect and Peierls methods) sets hoppings :math:`t` that connect equal
56 orbitals and equal spins of the two sites (:math:`t\,\mathbb{1}`), and
57 *set_onsite* with plain numbers shifts every orbital of a site. Onsite
58 energies can also be given orbital by orbital, and full hopping blocks
59 come from *set_slater_koster*, *set_spin_orbit*, *set_rashba* or
60 *set_hopping_block*. Call *set_peierls_phase* / *set_magnetic_field*
61 last: they multiply every hopping (and block) set so far.
63 Examples
64 --------
65 Example usage::
67 # graphene's sp3 bands: four orbitals per carbon
68 sys = OrbitalSystem(lat, orbitals={'a': ['s', 'px', 'py', 'pz'],
69 'b': ['s', 'px', 'py', 'pz']})
70 sys.set_onsite({'a': {'s': -8.87}, 'b': {'s': -8.87}})
71 sys.set_slater_koster(1, {'ss_sigma': -6.77, 'sp_sigma': 5.58,
72 'pp_sigma': 5.04, 'pp_pi': -3.03})
73 sys.get_ham()
74 '''
76 def __init__(self, lat: Lattice, orbitals: dict[str, list[str]] | None = None,
77 spin: bool = False, periodic: bool | tuple[bool, ...] = False) -> None:
78 System.__init__(self, lat, periodic)
79 error_handling.boolean(spin, 'spin')
80 if orbitals is None:
81 orbitals = {tag: ['s'] for tag in lat.tags}
82 error_handling.orbital_dict(orbitals, lat.tags)
83 self.orbitals = {tag: list(orbs) for tag, orbs in orbitals.items()}
84 self.spin = spin
85 self.ns = 2 if spin else 1
86 self.onsite_orb: dict = {} # tag -> {orbital: energy}
87 self.blocks: dict = {} # name -> list of (i, j, matrix): bond blocks, upper part
88 self.overlap_blocks: dict = {} # name -> list of (i, j, matrix)
89 self.onsite_blocks: dict = {} # name -> {site: matrix}
90 self.overlap = None
91 self._index()
93 # ------------------------------------------------------------------
94 # indexing
95 # ------------------------------------------------------------------
97 def _index(self) -> None:
98 '''
99 Private method. Rows of every site, orbital and spin.
100 '''
101 tags = self.lat.coor['tag']
102 self.norb_site = np.array([len(self.orbitals[t]) * self.ns for t in tags], dtype=int)
103 self.offset = np.concatenate([[0], np.cumsum(self.norb_site)[:-1]]).astype(int)
104 self.n_rows = int(self.norb_site.sum())
105 self.orb_site = np.repeat(np.arange(self.lat.sites), self.norb_site)
107 def row(self, site: int, orbital: str, spin: int = 0) -> int:
108 '''
109 Get the row of an orbital.
111 Parameters
112 ----------
113 site : Integer
114 Site index.
115 orbital : String
116 Orbital name (one of the site's).
117 spin : 0 (up) or 1 (down), default 0
119 Returns
120 -------
121 row : Integer
122 '''
123 self._index()
124 error_handling.site_index(site, self.lat.sites)
125 orbs = self.orbitals[self.lat.coor['tag'][site]]
126 error_handling.orbital_of_site(orbital, orbs)
127 error_handling.spin_index(spin, self.ns)
128 return int(self.offset[site] + orbs.index(orbital) * self.ns + spin)
130 def _site_orbs(self, site: int) -> list[str]:
131 return self.orbitals[self.lat.coor['tag'][site]]
133 def _kron_spin(self, mat: NDArray) -> NDArray:
134 '''
135 Private method. Orbital matrix times the spin identity.
136 '''
137 if not self.spin:
138 return np.asarray(mat, dtype='c16')
139 return np.kron(mat, np.eye(2))
141 def _same_orbitals(self, i: int, j: int) -> NDArray:
142 '''
143 Private method. Matrix connecting equal orbitals of sites i and j.
144 '''
145 oi, oj = self._site_orbs(i), self._site_orbs(j)
146 return np.array([[1. if a == b else 0. for b in oj] for a in oi])
148 # ------------------------------------------------------------------
149 # onsite terms
150 # ------------------------------------------------------------------
152 def set_onsite(self, dict_onsite: dict) -> None:
153 '''
154 Set onsite energies.
156 Parameters
157 ----------
158 dict_onsite : Dictionary
159 key: sublattice tag, val: onsite
160 energy (a number, applied to every orbital of the sites), or a
161 dictionary {orbital: energy} (orbitals not given get 0).
163 Examples
164 --------
165 Example usage::
167 sys.set_onsite({'a': 0.5, 'b': {'s': -8.87, 'pz': 0.}})
168 '''
169 error_handling.sites(self.lat.sites)
170 error_handling.set_onsite_orb(dict_onsite, self.lat.tags, self.orbitals)
171 numbers = {tag: val for tag, val in dict_onsite.items() if not isinstance(val, dict)}
172 System.set_onsite(self, numbers)
173 self.onsite_orb = {tag: dict(val) for tag, val in dict_onsite.items() if isinstance(val, dict)}
175 def set_zeeman(self, b: tuple[float, float, float]) -> None:
176 r'''
177 Set a Zeeman field :math:`\mathbf{b}\cdot\boldsymbol\sigma` on every
178 orbital (``spin=True`` only).
180 Parameters
181 ----------
182 b : Tuple of three real numbers :math:`(b_x, b_y, b_z)`
183 '''
184 error_handling.spinful(self.spin)
185 error_handling.vector3(b, 'b')
186 field = sum(c * s for c, s in zip(b, SIGMA))
187 self.onsite_blocks['zeeman'] = {
188 i: np.kron(np.eye(len(self._site_orbs(i))), field) for i in range(self.lat.sites)}
190 def set_atomic_soc(self, lam: float, tags: list[str] | None = None) -> None:
191 r'''
192 Set the atomic spin-orbit coupling
193 :math:`\lambda\,\mathbf{L}\cdot\mathbf{S} = \frac{\lambda}{2}\sum_k L_k\sigma_k`
194 on the p orbitals (``spin=True`` only; the sites must carry all three
195 of px, py, pz). It splits them into :math:`j = 3/2` (at
196 :math:`+\lambda/2`) and :math:`j = 1/2` (at :math:`-\lambda`) levels.
198 Parameters
199 ----------
200 lam : Real number
201 :math:`\lambda`.
202 tags : List of tags
203 Default value None (every tag with p orbitals).
204 '''
205 error_handling.spinful(self.spin)
206 error_handling.real_number(lam, 'lam')
207 if tags is None:
208 tags = [t for t, orbs in self.orbitals.items() if 'px' in orbs]
209 error_handling.p_shell(tags, self.orbitals)
210 L = orbital_angular_momentum()
211 blocks = {}
212 for i in range(self.lat.sites):
213 orbs = self._site_orbs(i)
214 if self.lat.coor['tag'][i] not in tags:
215 continue
216 mat = np.zeros((2 * len(orbs), 2 * len(orbs)), 'c16')
217 idx = [orbs.index(o) for o in ('px', 'py', 'pz')]
218 for k in range(3):
219 emb = np.zeros((len(orbs), len(orbs)), 'c16')
220 emb[np.ix_(idx, idx)] = L[k]
221 mat += 0.5 * lam * np.kron(emb, SIGMA[k])
222 blocks[i] = mat
223 self.onsite_blocks['atomic_soc'] = blocks
225 # ------------------------------------------------------------------
226 # bond blocks
227 # ------------------------------------------------------------------
229 def _bonds(self, n: int) -> tuple[NDArray, NDArray, NDArray]:
230 '''
231 Private method. Bonds of order n (oriented, see *get_bonds*) and their vectors.
232 '''
233 self.get_distances(n)
234 self.nmax = len(self.dist_uni) - 1
235 error_handling.positive_int_lim(n, 'n', self.nmax)
236 i, j, _ = self.get_bonds(n)
237 return i, j, self._bond_vectors(i, j)
239 def set_slater_koster(self, n: int, params: dict, overlap: dict | None = None) -> None:
240 r'''
241 Set the hoppings between the orbitals of every pair of *n*-th
242 neighbours from Slater-Koster bond integrals (see
243 :func:`tbkit.slater_koster.sk_block`), replacing those set before for
244 the same *n*. The blocks are spin independent.
246 Parameters
247 ----------
248 n : Positive integer
249 Neighbour order (1 for nearest neighbours).
250 params
251 Dictionary of bond integrals (e.g. ``{'pp_sigma': 5.04,
252 'pp_pi': -3.03}``), or dictionary {tag pair: bond integrals} (e.g.
253 ``{'ab': {...}, 'ba': {...}}``, for bonds from an 'a' to a 'b' site).
254 overlap : Dictionary, default None
255 Bond integrals of the
256 overlap matrix :math:`S_{ij} = \langle i|j\rangle`, in the same
257 format: the basis is then non-orthogonal, and *get_eig* solves
258 :math:`H\psi = ES\psi`.
259 '''
260 error_handling.positive_int(n, 'n')
261 i, j, d = self._bonds(n)
262 tags = self.lat.coor['tag']
263 for store, par in ((self.blocks, params), (self.overlap_blocks, overlap)):
264 if par is None:
265 store.pop(('sk', n), None)
266 continue
267 by_pair = error_handling.sk_pair_params(par, tags[i], tags[j])
268 blocks = []
269 for a, b, vec in zip(i, j, d):
270 p = par[tags[a] + tags[b]] if by_pair else par
271 blocks.append((a, b, self._kron_spin(sk_block(self._site_orbs(a), self._site_orbs(b), vec, p))))
272 store[('sk', n)] = blocks
274 def set_spin_orbit(self, lam: float, n: int = 2) -> None:
275 r'''
276 Set the Kane-Mele intrinsic spin-orbit coupling on the *n*-th
277 neighbours (``spin=True`` only),
279 .. math::
281 H_{SO} = i\lambda\sum_{\langle\langle ij\rangle\rangle}\nu_{ij}\,
282 c_i^\dagger\sigma_z c_j\, ,\qquad
283 \nu_{ij} = \sum_k \mathrm{sign}\left[(\mathbf{d}_{ik}\times\mathbf{d}_{kj})_z\right]\, ,
285 the sum running over the nearest neighbours :math:`k` shared by
286 :math:`i` and :math:`j` (one on the honeycomb lattice, where
287 :math:`\nu_{ij} = \pm1` says whether the path turns left or right).
288 It connects equal orbitals of the two sites.
290 Parameters
291 ----------
292 lam : Real number
293 :math:`\lambda`.
294 n : Positive integer, default 2
295 Neighbour order.
296 '''
297 error_handling.spinful(self.spin)
298 error_handling.real_number(lam, 'lam')
299 error_handling.positive_int(n, 'n')
300 ni, nj, _ = self._bonds(1)
301 neighbours: list[set[int]] = [set() for _ in range(self.lat.sites)]
302 for a, b in zip(ni, nj):
303 neighbours[a].add(b)
304 neighbours[b].add(a)
305 i, j, _ = self._bonds(n)
306 coords = self._coords()
307 blocks = []
308 for a, b in zip(i, j):
309 nu = 0.
310 for k in neighbours[a] & neighbours[b]:
311 d1, d2 = self._min_image(coords[k] - coords[a]), self._min_image(coords[b] - coords[k])
312 nu += np.sign(d1[0] * d2[1] - d1[1] * d2[0])
313 if nu:
314 blocks.append((a, b, 1j * lam * nu * np.kron(self._same_orbitals(a, b), SIGMA[2])))
315 self.blocks[('soc', n)] = blocks
317 def set_rashba(self, lam: float, n: int = 1) -> None:
318 r'''
319 Set a Rashba spin-orbit coupling on the *n*-th neighbours
320 (``spin=True`` only),
321 :math:`H_R = i\lambda\sum_{\langle ij\rangle} c_i^\dagger
322 (\boldsymbol\sigma\times\hat{\mathbf{d}}_{ij})_z\, c_j`,
323 :math:`\hat{\mathbf{d}}_{ij}` the bond direction. It connects equal
324 orbitals of the two sites.
326 Parameters
327 ----------
328 lam : Real number
329 :math:`\lambda`.
330 n : Positive integer, default 1
331 '''
332 error_handling.spinful(self.spin)
333 error_handling.real_number(lam, 'lam')
334 error_handling.positive_int(n, 'n')
335 i, j, d = self._bonds(n)
336 blocks = []
337 for a, b, vec in zip(i, j, d):
338 u = vec / np.linalg.norm(vec)
339 spin = 1j * lam * (SIGMA[0] * u[1] - SIGMA[1] * u[0])
340 blocks.append((a, b, np.kron(self._same_orbitals(a, b), spin)))
341 self.blocks[('rashba', n)] = blocks
343 def set_hopping_block(self, dict_block: dict[tuple[int, int], NDArray]) -> None:
344 r'''
345 Set hopping blocks by hand, added to the others: key: site pair
346 :math:`(i, j)`, val: matrix of shape (rows of i, rows of j) --
347 orbitals, then spin, as *row* orders them. Their Hermitian conjugates
348 are added.
350 Parameters
351 ----------
352 dict_block : Dictionary
353 '''
354 error_handling.set_hopping_def(None, {k: 0. for k in dict_block}, self.lat.sites, 'dict_block')
355 self._index()
356 blocks = self.blocks.setdefault(('manual', 0), [])
357 for (a, b), mat in dict_block.items():
358 mat = np.asarray(mat, dtype='c16')
359 error_handling.block_shape(mat, (self.norb_site[a], self.norb_site[b]))
360 blocks.append((a, b, mat))
362 def clear_hopping(self) -> None:
363 '''
364 Clear the hoppings: *hop*, and every hopping block.
365 '''
366 System.clear_hopping(self)
367 self.blocks = {}
368 self.overlap_blocks = {}
370 def set_peierls_phase(self, phase: Callable[..., NDArray]) -> None:
371 r'''
372 Apply the Peierls substitution (see *System.set_peierls_phase*) to
373 the hoppings *hop* and to every hopping block set so far.
375 Parameters
376 ----------
377 phase : Callable ``phase(xi, yi, xj, yj)``
378 '''
379 error_handling.is_callable(phase, 'phase')
380 if self.hop.size:
381 System.set_peierls_phase(self, phase)
382 x, y = self.lat.coor['x'], self.lat.coor['y']
383 for key, blocks in self.blocks.items():
384 # the end of the short bond (outside the sample if it wraps around)
385 d = self._bond_vectors(np.array([a for a, _, _ in blocks], int),
386 np.array([b for _, b, _ in blocks], int))
387 self.blocks[key] = [(a, b, mat * np.exp(1j * phase(x[a], y[a], x[a] + dx, y[a] + dy)))
388 for (a, b, mat), (dx, dy) in zip(blocks, d[:, :2])]
390 # ------------------------------------------------------------------
391 # Hamiltonian and eigenstates
392 # ------------------------------------------------------------------
394 def _assemble(self, blocks: list, n: int) -> sparse.csr_matrix:
395 '''
396 Private method. Sparse matrix of bond blocks, upper part only.
397 '''
398 rows, cols, vals = [], [], []
399 for a, b, mat in blocks:
400 r = self.offset[a] + np.arange(mat.shape[0])
401 c = self.offset[b] + np.arange(mat.shape[1])
402 rr, cc = np.meshgrid(r, c, indexing='ij')
403 rows.append(rr.ravel())
404 cols.append(cc.ravel())
405 vals.append(mat.ravel())
406 if not rows:
407 return sparse.csr_matrix((n, n), dtype='c16')
408 return sparse.csr_matrix((np.concatenate(vals), (np.concatenate(rows), np.concatenate(cols))),
409 shape=(n, n))
411 def get_ham(self, **params) -> None:
412 '''
413 Get the Hamiltonian (and, for a non-orthogonal basis, the overlap
414 matrix *overlap*).
416 Parameters
417 ----------
418 params
419 Values of the parameters of the value functions of
420 *hop* (see *System.get_ham*).
421 '''
422 error_handling.orbital_terms(self.hop, self.blocks)
423 self._index()
424 n = self.n_rows
425 ham = sparse.csr_matrix((n, n), dtype='c16')
426 if self.hop.size:
427 error_handling.hop_sites(self.hop, self.lat.sites)
428 rows, cols, vals = [], [], []
429 names = sorted({o for orbs in self.orbitals.values() for o in orbs})
430 index = {o: np.full(self.lat.sites, -1) for o in names}
431 for site in range(self.lat.sites):
432 for k, o in enumerate(self._site_orbs(site)):
433 index[o][site] = self.offset[site] + k * self.ns
434 hi, hj = self.hop['i'].astype(int), self.hop['j'].astype(int)
435 t = self._hop_t(params)
436 for o in names:
437 keep = (index[o][hi] >= 0) & (index[o][hj] >= 0)
438 for s in range(self.ns):
439 rows.append(index[o][hi[keep]] + s)
440 cols.append(index[o][hj[keep]] + s)
441 vals.append(t[keep])
442 scalar = sparse.csr_matrix((np.concatenate(vals), (np.concatenate(rows), np.concatenate(cols))),
443 shape=(n, n))
444 if np.all(self.hop['ang'] >= 0) or np.all(self.hop['ang'] < 0):
445 scalar = scalar + scalar.conj().T
446 ham = ham + scalar
447 bonds = self._assemble([blk for blocks in self.blocks.values() for blk in blocks], n)
448 ham = ham + bonds + bonds.conj().T
449 diag = np.zeros(n, 'c16')
450 if self.onsite.size == self.lat.sites:
451 diag += self.onsite[self.orb_site]
452 for tag, energies in self.onsite_orb.items():
453 for site in np.flatnonzero(self.lat.coor['tag'] == tag).tolist():
454 for o, e in energies.items():
455 k = self._site_orbs(site).index(o)
456 diag[self.offset[site] + k * self.ns: self.offset[site] + (k + 1) * self.ns] += e
457 ham = ham + sparse.diags(diag)
458 for blocks in self.onsite_blocks.values():
459 ham = ham + self._assemble([(i, i, mat) for i, mat in blocks.items()], n)
460 self.ham = sparse.csr_matrix(ham)
461 self.overlap = None
462 overlap = [blk for blocks in self.overlap_blocks.values() for blk in blocks]
463 if overlap:
464 o = self._assemble(overlap, n)
465 self.overlap = sparse.csr_matrix(sparse.identity(n, dtype='c16') + o + o.conj().T)
467 def _set_states(self) -> None:
468 '''
469 Private method. Intensities summed over each site's orbitals and spin
470 (Mulliken populations for a non-orthogonal basis), and sublattice
471 polarizations.
472 '''
473 if self.overlap is None:
474 weight = np.abs(self.rn) ** 2
475 else:
476 weight = np.real(self.rn.conj() * (self.overlap @ self.rn))
477 self.intensity = np.zeros((self.lat.sites, self.rn.shape[1]))
478 np.add.at(self.intensity, self.orb_site, weight)
479 self.pola = np.zeros((self.rn.shape[1], len(self.lat.tags)))
480 for i, tag in enumerate(self.lat.tags):
481 self.pola[:, i] = np.sum(self.intensity[self.lat.coor['tag'] == tag, :], axis=0)
483 def get_eig(self, eigenvec: bool = False, left: bool = False) -> None:
484 r'''
485 Get the eigenenergies (and eigenvectors), as *System.get_eig*. With a
486 non-orthogonal basis (see *set_slater_koster*), solve the generalized
487 problem :math:`H\psi = ES\psi` (eigenvectors normalized as
488 :math:`\psi^\dagger S\psi = 1` for a Hermitian :math:`H`).
490 Parameters
491 ----------
492 eigenvec : Boolean, default False
493 left : Boolean
494 Default value False (orthogonal basis only).
495 '''
496 if self.overlap is None:
497 return System.get_eig(self, eigenvec, left)
498 error_handling.empty_ham(self.ham)
499 error_handling.boolean(eigenvec, 'eigenvec')
500 error_handling.no_left(left)
501 self.ln = np.array([], 'c16')
502 ham, s = self.ham.toarray(), self.overlap.toarray()
503 if (self.ham.conj().T != self.ham).nnz:
504 en, vec = LA.eig(ham, s)
505 ind = np.argsort(en.real, kind='stable')
506 self.en, self.rn = en[ind], vec[:, ind]
507 else:
508 self.en, self.rn = LA.eigh(ham, s)
509 if eigenvec:
510 self._set_states()
512 def _row_sites(self) -> NDArray:
513 '''
514 Private method. Site of every row of the Hamiltonian.
515 '''
516 self._index()
517 return self.orb_site
519 def get_eig_sparse(self, n_eig: int = 10, sigma: complex = 0., eigenvec: bool = False) -> None:
520 '''
521 Get the eigenpairs closest to *sigma* (see *System.get_eig_sparse*;
522 orthogonal basis only).
523 '''
524 error_handling.orthogonal(self.overlap)
525 System.get_eig_sparse(self, n_eig, sigma, eigenvec)
527 def get_local_chern_marker(self, e_fermi: float = 0., area: float | None = None) -> NDArray[np.float64]:
528 '''
529 Get the local Chern marker (see *System.get_local_chern_marker*),
530 summed over the orbitals and spin of each site (orthogonal basis only).
532 Parameters
533 ----------
534 e_fermi : Real number, default 0
535 area : Positive real number, default None
537 Returns
538 -------
539 marker : Real ndarray, shape (sites,)
540 '''
541 error_handling.orthogonal(self.overlap)
542 return System.get_local_chern_marker(self, e_fermi, area)
544 def get_bott_index(self, e_fermi: float = 0.) -> float:
545 '''
546 Get the Bott index of a torus (see *System.get_bott_index*), the
547 orbitals and spin of a site sharing its position (orthogonal basis only).
549 Parameters
550 ----------
551 e_fermi : Real number, default 0
553 Returns
554 -------
555 bott : Real number, close to an integer
556 '''
557 error_handling.orthogonal(self.overlap)
558 return System.get_bott_index(self, e_fermi)