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