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