Coverage for tbkit/bridges.py: 100%
122 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 11:04 +0000
1r'''
2Bridges between the two pipelines of tbkit: the real-space *System* (a
3finite lattice, its sparse Hamiltonian) and the reciprocal-space *KSpace*
4(the Bloch Hamiltonian :math:`H(\mathbf{k})` of the infinite lattice).
6* *kspace_from_system* reads the Bloch model off a translation-invariant
7 *System* built the real-space way (*get_lattice*, *set_hopping* by
8 neighbour order, *set_onsite*, ...).
9* *finite_model* and *finite_system* cut a finite sample (open boundaries,
10 or a torus) out of a *KSpace* model: its sparse Hamiltonian with the site
11 positions, or a ready-to-use *System*.
13Both use the convention of *KSpace.finite_ham*: the hopping
14:math:`t_{ij}(\mathbf{R}) = \langle i,\mathbf{0}|H|j,\mathbf{R}\rangle`
15joins orbital *i* of every cell to orbital *j* of the cell :math:`\mathbf{R}`
16further. So, for the same model, the spectrum of a System on an
17:math:`N_1\times N_2` torus is the set of KSpace bands on the
18:math:`N_1\times N_2` Brillouin-zone mesh.
19'''
20from __future__ import annotations
22import numpy as np
23from numpy.typing import NDArray
24import scipy.sparse as sp
26import tbkit.error_handling as error_handling
27import tbkit.neighbours as neighbours
28from tbkit.kspace import KSpace
29from tbkit.lattice import Lattice
30from tbkit.orbital import OrbitalSystem
31from tbkit.system import System, HOP_DTYPE, ATOL, _upper
34def _coords(lat: Lattice) -> NDArray[np.float64]:
35 '''
36 Private function. Site coordinates of a finite lattice, shape (sites, space_dim).
37 '''
38 return np.stack([lat.coor[f] for f in ('x', 'y', 'z')[:lat.space_dim]], axis=1)
41def cell_orbitals(lat: Lattice) -> tuple[NDArray[np.int64], NDArray[np.int64]]:
42 r'''
43 Get, for every site of a finite lattice (*lat.coor*), the unit-cell
44 orbital it is a copy of and its cell: site *s* sits at
45 :math:`\boldsymbol\tau_{o_s} + \sum_k n_{s,k}\mathbf{a}_k`, with the tag
46 of orbital :math:`o_s` (within ``ATOL``).
48 Parameters
49 ----------
50 lat
51 **Lattice** instance, after *get_lattice* (sites may have
52 been removed, but not moved).
54 Returns
55 -------
56 cells : Integer ndarray, shape (sites, dim)
57 :math:`n_{s,k}`.
58 orbitals : Integer ndarray, shape (sites,)
59 :math:`o_s`.
60 '''
61 error_handling.lat(lat)
62 error_handling.empty_coor(lat.coor)
63 a, tau, tags = neighbours._geometry(lat)
64 pos = _coords(lat)
65 pinv = np.linalg.pinv(a)
66 match = np.zeros((lat.sites, len(tau)), bool)
67 cells = np.zeros((lat.sites, len(tau), len(a)), int)
68 for o in range(len(tau)):
69 rel = pos - tau[o]
70 n = np.rint(rel @ pinv).astype(int)
71 match[:, o] = (np.linalg.norm(n @ a - rel, axis=1) < ATOL) & (lat.coor['tag'] == tags[o])
72 cells[:, o] = n
73 error_handling.lattice_sites(match.sum(axis=1))
74 orb = np.argmax(match, axis=1)
75 return cells[np.arange(lat.sites), orb], orb
78def kspace_from_system(sys: System, periodic: bool | None = None, tol: float = 1e-9) -> KSpace:
79 r'''
80 Get the Bloch Hamiltonian of a periodic model defined the real-space
81 way: every matrix element :math:`H_{ab}` of the finite *System* (built
82 by *get_ham*) joins orbital :math:`o_a` of cell :math:`\mathbf{n}_a` to
83 orbital :math:`o_b` of cell :math:`\mathbf{n}_b` (see *cell_orbitals*),
84 i.e. is the hopping :math:`t_{o_ao_b}(\mathbf{n}_b-\mathbf{n}_a)` of the
85 infinite lattice. Each such hopping must take the same value in every
86 cell where it appears (bonds cut by the boundary are simply missing).
87 The lattice must be large enough to hold every bond of the model.
89 A System on a torus (*periodic*, e.g. from *finite_system*) has bonds
90 that wrap around: their lattice vectors are then taken modulo the
91 :math:`N_1\times N_2\times\dots` cells of *get_lattice*, as the shortest
92 image (so the model's range must stay below half the torus).
94 A Hermitian hopping matrix gives a Hermitian model (one representative
95 per bond, *set_hopping* adding the conjugate); otherwise every matrix
96 element is kept, with ``hermitian=False``. Onsite energies are taken
97 orbital by orbital.
99 Parameters
100 ----------
101 sys : **System** instance with its hoppings set (one orbital per site)
102 periodic : Boolean
103 Default value None: True if *sys* has
104 periodic boundaries along some primitive vector (``System(lat,
105 periodic=...)``, or *finite_system* with ``periodic=True``). The
106 System is a torus of the *lat.n1* x *lat.n2* (x *lat.n3*) cells of
107 *get_lattice*.
108 tol : Positive real number, default 1e-9
109 Largest allowed
110 difference between equivalent matrix elements.
112 Returns
113 -------
114 ks
115 **KSpace** instance (spinless), on a new **Lattice**
116 with the *unit_cell* and *prim_vec* of *sys.lat*.
118 Examples
119 --------
120 Example usage::
122 lat = lattices.honeycomb()
123 lat.get_lattice(6, 6)
124 sys = System(lat)
125 sys.set_hopping([{'n': 1, 't': -1.}, {'n': 2, 't': 0.1}])
126 gra = kspace_from_system(sys) # the Bloch model of the same graphene
127 '''
128 error_handling.sys(sys)
129 error_handling.not_orbital_system(sys, OrbitalSystem)
130 if periodic is None:
131 periodic = any(sys.periodic)
132 error_handling.boolean(periodic, 'periodic')
133 error_handling.positive_real(tol, 'tol')
134 sys.get_ham()
135 cells, orb = cell_orbitals(sys.lat)
136 ham = sys.ham.tocoo()
137 rows, cols, vals = ham.row, ham.col, ham.data.astype('c16')
138 R = cells[cols] - cells[rows]
139 if periodic:
140 sizes = np.array((sys.lat.n1, sys.lat.n2, sys.lat.n3)[:R.shape[1]])
141 R = (R + sizes // 2) % sizes - sizes // 2 # shortest image, in [-N/2, N/2)
142 error_handling.torus_range(R, sizes)
143 keys = np.column_stack([orb[rows], orb[cols], R])
144 uniq, first, inv = np.unique(keys, axis=0, return_index=True, return_inverse=True)
145 inv = inv.ravel()
146 values = vals[first]
147 error_handling.translation_invariant(float(np.abs(vals - values[inv]).max()), tol)
148 off = ham - sp.diags(ham.diagonal())
149 hermitian = bool(abs(off - off.conj().T).max() <= tol)
150 lat = Lattice(unit_cell=list(sys.lat.unit_cell), prim_vec=list(sys.lat.prim_vec))
151 ks = KSpace(lat)
152 list_hop = []
153 for key, t in zip(uniq, values):
154 i, j, R = int(key[0]), int(key[1]), tuple(int(n) for n in key[2:])
155 if i == j and not any(R):
156 ks.onsite[i] = t
157 continue
158 nonzero = [n for n in R if n != 0]
159 if hermitian and not ((nonzero and nonzero[0] > 0) or (not nonzero and i < j)):
160 continue # the conjugate of a kept representative
161 list_hop.append({'i': i, 'j': j, 'R': R, 't': complex(t)})
162 ks.set_hopping(list_hop, hermitian=hermitian)
163 return ks
166def _finite_sites(ks: KSpace, n_cells: tuple[int, ...]) -> tuple[NDArray[np.float64], NDArray]:
167 '''
168 Private function. Positions and tags of the rows of *ks.finite_ham(n_cells)*.
169 '''
170 n_tot = int(np.prod(n_cells))
171 cells = np.array(np.unravel_index(np.arange(n_tot), n_cells, order='F')).T
172 a = np.array(ks.lat.prim_vec, dtype='f8')
173 pos = (cells @ a)[:, None, :] + ks.orbital_positions()[None, :, :]
174 tags = np.repeat(ks.tags, 2) if ks.spin else ks.tags
175 return pos.reshape(n_tot * ks.norb, ks.space_dim), np.tile(tags, n_tot)
178def finite_model(
179 ks: KSpace, n_cells: int | tuple[int, ...], periodic: bool = False,
180) -> tuple[sp.csr_matrix, NDArray[np.float64], NDArray]:
181 r'''
182 Get a finite sample of a *KSpace* model: its sparse Hamiltonian (that of
183 *KSpace.finite_ham*, open or periodic boundaries), with the position
184 and tag of every row. Spinful models are supported (the two spin
185 components of a site share its position).
187 Parameters
188 ----------
189 ks : **KSpace** instance
190 n_cells : Positive integer, or tuple of *dim* positive integers
191 Number of unit cells along each primitive vector.
192 periodic : Boolean, default False
193 Periodic boundaries (a torus).
195 Returns
196 -------
197 ham : Complex CSR matrix, shape (N*norb, N*norb)
198 positions : Real ndarray, shape (N*norb, space_dim)
199 Row
200 :math:`c\,\mathrm{norb}+o` sits at :math:`\boldsymbol\tau_o + \sum_k n_k\mathbf{a}_k`
201 (cell *c* has :math:`c = n_1 + N_1 n_2 + N_1N_2 n_3`).
202 tags : String ndarray, shape (N*norb,)
203 Sublattice tags.
204 '''
205 error_handling.kspace(ks, KSpace)
206 error_handling.nk(n_cells, ks.dim)
207 error_handling.boolean(periodic, 'periodic')
208 n_cells = (n_cells,) * ks.dim if isinstance(n_cells, int) else n_cells
209 pos, tags = _finite_sites(ks, n_cells)
210 return ks.finite_ham(n_cells, periodic=periodic, sparse=True), pos, tags
213def finite_system(ks: KSpace, n_cells: int | tuple[int, ...], periodic: bool = False) -> System:
214 r'''
215 Build a finite *System* out of a (spinless) *KSpace* model: *lat.coor*
216 holds the sites of :math:`N_1\times N_2\times\dots` unit cells (in the
217 order of *finite_model*), and *sys.hop* / *sys.onsite* the model's
218 hoppings and onsite energies, so that *sys.get_ham()* gives
219 *ks.finite_ham(n_cells, periodic)*. Every real-space tool (*get_eig*,
220 *get_ldos*, *Plot*, disorder, ...) then applies.
222 In *sys.hop*, each bond is stored with the conventions of
223 *System.set_hopping*: a Hermitian model keeps one orientation per bond
224 (angle in [0, 180), *get_ham* adds the conjugate); a non-reciprocal
225 one keeps both, the missing direction with a zero amplitude. 'n' is the
226 neighbour order of the bond (0 for coinciding orbitals), 'tag' the
227 sublattice pair. On a torus, a bond that wraps around keeps its short
228 bond vector for its angle, and a hopping that wraps onto its own site
229 is added to the onsite energy; the System has periodic boundaries
230 (``sys.periodic``), so that further *set_hopping* calls, *Plot* and
231 *get_bott_index* treat it as a torus.
233 Parameters
234 ----------
235 ks : **KSpace** instance, with ``spin=False``
236 n_cells : Positive integer, or tuple of *dim* positive integers
237 periodic : Boolean, default False
238 Periodic boundaries (a torus).
240 Returns
241 -------
242 sys : **System** instance, with its hoppings set
244 Examples
245 --------
246 Example usage::
248 sys = finite_system(gra, (30, 30), periodic=True)
249 sys.get_ham()
250 sys.get_eig() # = gra.mesh_bands((30, 30)), sorted
251 '''
252 error_handling.kspace(ks, KSpace)
253 error_handling.nk(n_cells, ks.dim)
254 error_handling.boolean(periodic, 'periodic')
255 error_handling.spinless(ks.spin)
256 n_cells = (n_cells,) * ks.dim if isinstance(n_cells, int) else n_cells
257 pos, tags = _finite_sites(ks, n_cells)
258 lat = Lattice(unit_cell=list(ks.lat.unit_cell), prim_vec=list(ks.lat.prim_vec))
259 coor = np.zeros(len(pos), dtype=lat.dtype)
260 for c, f in enumerate(('x', 'y', 'z')[:ks.space_dim]):
261 coor[f] = pos[:, c]
262 coor['tag'] = tags
263 lat.coor, lat.sites = coor, len(coor)
264 lat.n1, lat.n2, lat.n3 = (tuple(n_cells) + (1, 1))[:3]
265 sys = System(lat, periodic=periodic)
266 _, n_tot, entries = ks._finite_entries(n_cells, periodic)
267 tau = ks.orbital_positions()
268 rows = np.concatenate([np.zeros(0, int)] + [e[0] for e in entries])
269 cols = np.concatenate([np.zeros(0, int)] + [e[1] for e in entries])
270 t = np.concatenate([np.zeros(0, 'c16')] + [np.full(len(e[0]), e[6], 'c16') for e in entries])
271 d = np.concatenate([np.zeros((0, ks.space_dim))]
272 + [np.repeat((e[5] + tau[e[3]] - tau[e[2]])[None], len(e[0]), axis=0)
273 for e in entries])
274 n_sites = len(pos)
275 diag = rows == cols
276 sys.onsite = np.tile(ks.onsite, n_tot) + (np.bincount(rows[diag], t[diag].real, n_sites)
277 + 1j * np.bincount(rows[diag], t[diag].imag, n_sites))
278 rows, cols, t, d = rows[~diag], cols[~diag], t[~diag], d[~diag]
279 dist = np.linalg.norm(d, axis=1)
280 zero = dist <= ATOL
281 ang = neighbours._bond_angles(d)
282 up = np.where(zero, rows < cols, _upper(ang))
283 ang_up = np.where(zero, 0., np.maximum(np.where(up, ang, neighbours._bond_angles(-d)), 0.))
284 if ks._nonreciprocal:
285 i = np.concatenate([rows, cols])
286 j = np.concatenate([cols, rows])
287 t = np.concatenate([t, np.zeros(len(t), 'c16')])
288 ang = np.concatenate([np.where(up, ang_up, ang_up - 180.), np.where(up, ang_up - 180., ang_up)])
289 dist = np.concatenate([dist, dist])
290 else:
291 i, j, t, ang, dist = rows[up], cols[up], t[up], ang_up[up], dist[up]
292 hop = np.zeros(len(i), dtype=HOP_DTYPE)
293 hop['n'] = neighbours.shell_index(ks.lat, dist)
294 hop['i'], hop['j'], hop['t'], hop['ang'] = i, j, t, ang
295 hop['tag'] = np.char.add(tags[i], tags[j])
296 sys.hop = hop
297 return sys