Coverage for tbkit/io.py: 100%
154 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 11:04 +0000
1r'''
2Save and load models, and import models from other codes.
4*save_model* / *load_model* store a **Lattice**, **System** or **KSpace**
5in a versioned NumPy ``.npz`` archive (no pickle, readable with NumPy
6alone): the lattice (*unit_cell*, *prim_vec*, and the sites of
7*get_lattice* if any), and for a System its onsite energies and hopping
8array *sys.hop*, for a KSpace its onsite terms (with the spin-off-diagonal
9block), hoppings, overlaps and *spin*. Every number is stored in binary,
10so a loaded model reproduces the Hamiltonians exactly.
12Archive layout (format ``'tbkit-model'``, version :data:`VERSION`):
14* ``format``, ``version``, ``kind`` (``'Lattice'``, ``'System'`` or
15 ``'KSpace'``), ``tbkit_version``;
16* ``uc_tag`` (n_sites,), ``uc_r0`` (n_sites, space_dim), ``prim_vec``
17 (dim, space_dim), ``coor`` (the structured site array), ``n_cells`` (3,);
18* System: ``onsite`` (sites,), ``hop`` (the structured *sys.hop*), and,
19 from version 2, ``periodic`` (one boolean per primitive vector) for a
20 System with periodic boundaries (every other model is still written as
21 version 1, which older tbkit versions read);
22* KSpace: ``spin``, ``onsite`` (norb,), ``onsite_offdiag`` (norb, norb),
23 ``nonreciprocal``, and ``hop_i``, ``hop_j``, ``hop_R`` (Cartesian
24 lattice vectors), ``hop_t`` for the hoppings, ``ovl_i``, ``ovl_j``,
25 ``ovl_R``, ``ovl_t`` for the overlaps.
27*read_wannier90* imports a tight-binding model from Wannier90's
28``seedname_hr.dat`` (see there).
29'''
30from __future__ import annotations
32import os
33from typing import Any, cast
35import numpy as np
36from numpy.typing import ArrayLike, NDArray
38import tbkit
39import tbkit.error_handling as error_handling
40from tbkit.kspace import KSpace
41from tbkit.lattice import Lattice
42from tbkit.orbital import OrbitalSystem
43from tbkit.system import System
46#: Name of the archive format written by *save_model*.
47FORMAT = 'tbkit-model'
48#: Version of the archive format written by *save_model*; *load_model*
49#: reads this version and every earlier one.
50VERSION = 2
53def _hop_arrays(hops: list, space_dim: int, prefix: str) -> dict:
54 '''
55 Private function. A KSpace hopping list (i, j, R_cart, t) as arrays.
56 '''
57 return {prefix + '_i': np.array([h[0] for h in hops], dtype='i8'),
58 prefix + '_j': np.array([h[1] for h in hops], dtype='i8'),
59 prefix + '_R': np.array([h[2] for h in hops], dtype='f8').reshape(len(hops), space_dim),
60 prefix + '_t': np.array([h[3] for h in hops], dtype='c16')}
63def _hop_list(data, prefix: str) -> list:
64 '''
65 Private function. Inverse of *_hop_arrays*.
66 '''
67 return [(int(i), int(j), R.copy(), complex(t))
68 for i, j, R, t in zip(data[prefix + '_i'], data[prefix + '_j'],
69 data[prefix + '_R'], data[prefix + '_t'])]
72def save_model(model: Lattice | System | KSpace, path: str | os.PathLike) -> str:
73 '''
74 Save a **Lattice**, **System** or **KSpace** to a ``.npz`` archive (see
75 the module docstring for its layout). Subclasses are saved as their
76 base class (e.g. a *GrapheneSystem* as a *System*); an *OrbitalSystem*
77 or a driven (Floquet) model cannot be saved.
79 Parameters
80 ----------
81 model : **Lattice**, **System** or **KSpace** instance
82 path : String or path
83 File name; ``.npz`` is appended if missing.
85 Returns
86 -------
87 path : String
88 The file written.
90 Examples
91 --------
92 Example usage::
94 save_model(gra, 'graphene.npz')
95 gra2 = load_model('graphene.npz')
96 '''
97 error_handling.saveable_model(model, Lattice, System, KSpace, OrbitalSystem)
98 error_handling.no_value_functions(model, 'save_model')
99 error_handling.file_path(path, 'path')
100 path = os.fspath(path)
101 if not path.endswith('.npz'):
102 path += '.npz'
103 lat = model.lat if isinstance(model, (System, KSpace)) else model
104 kind = 'System' if isinstance(model, System) else 'KSpace' if isinstance(model, KSpace) else 'Lattice'
105 # only what needs version 2 (a periodic System) is written as such, so
106 # that older tbkit versions keep reading everything else
107 data: dict[str, Any] = {'format': np.array(FORMAT), 'version': np.array(1), 'kind': np.array(kind),
108 'tbkit_version': np.array(tbkit.__version__),
109 'uc_tag': np.array([dic['tag'] for dic in lat.unit_cell], dtype='U1'),
110 'uc_r0': np.array([dic['r0'] for dic in lat.unit_cell], dtype='f8'),
111 'prim_vec': np.array(lat.prim_vec, dtype='f8'),
112 'coor': lat.coor,
113 'n_cells': np.array([lat.n1, lat.n2, lat.n3], dtype='i8')}
114 if isinstance(model, System):
115 data['onsite'] = np.asarray(model.onsite, dtype='c16')
116 data['hop'] = model.hop
117 if any(model.periodic):
118 data['periodic'] = np.array(model.periodic, dtype=bool)
119 data['version'] = np.array(VERSION)
120 elif isinstance(model, KSpace):
121 data.update({'spin': np.array(model.spin), 'onsite': model.onsite,
122 'onsite_offdiag': model._onsite_offdiag,
123 'nonreciprocal': np.array(model._nonreciprocal)})
124 data.update(_hop_arrays(model._hop, model.space_dim, 'hop'))
125 data.update(_hop_arrays(model._overlap_hop, model.space_dim, 'ovl'))
126 np.savez(path, **data)
127 return path
130def load_model(path: str | os.PathLike) -> Lattice | System | KSpace:
131 '''
132 Load a model saved by *save_model*.
134 Parameters
135 ----------
136 path : String or path
137 File name.
139 Returns
140 -------
141 model
142 **Lattice**, **System** or **KSpace** instance, whose
143 Hamiltonians (*get_ham*) are exactly those of the saved model.
144 '''
145 error_handling.file_path(path, 'path')
146 with np.load(path, allow_pickle=False) as data:
147 error_handling.model_archive(data, FORMAT, VERSION)
148 unit_cell = [{'tag': str(tag), 'r0': tuple(float(c) for c in r0)}
149 for tag, r0 in zip(data['uc_tag'], data['uc_r0'])]
150 prim_vec = [tuple(float(c) for c in a) for a in data['prim_vec']]
151 lat = Lattice(unit_cell=unit_cell, prim_vec=prim_vec)
152 lat.coor = data['coor'].astype(lat.dtype)
153 lat.sites = len(lat.coor)
154 lat.n1, lat.n2, lat.n3 = (int(n) for n in data['n_cells'])
155 if lat.sites:
156 lat.tags = np.unique(np.concatenate([lat.tags, lat.coor['tag']]))
157 kind = str(data['kind'])
158 if kind == 'Lattice':
159 return lat
160 if kind == 'System':
161 periodic = tuple(bool(p) for p in data['periodic']) if 'periodic' in data.files else False
162 sys = System(lat, periodic=periodic)
163 sys.onsite = data['onsite'].copy()
164 sys.hop = data['hop'].copy()
165 return sys
166 ks = KSpace(lat, spin=bool(data['spin']))
167 ks.onsite = data['onsite'].copy()
168 ks._onsite_offdiag = data['onsite_offdiag'].copy()
169 ks._nonreciprocal = bool(data['nonreciprocal'])
170 ks._hop = _hop_list(data, 'hop')
171 ks._overlap_hop = _hop_list(data, 'ovl')
172 return ks
175#################################
176# WANNIER90
177#################################
179#: Bohr radius in Angstrom (CODATA 2018), for ``unit_cell_cart`` given in bohr.
180BOHR = 0.529177210903
183def read_hr(path: str | os.PathLike) -> tuple[int, NDArray[np.int64], NDArray[np.int64], NDArray[np.complex128]]:
184 r'''
185 Read a Wannier90 ``seedname_hr.dat`` file: a header line, the number of
186 Wannier functions, the number of lattice vectors :math:`\mathbf{R}`,
187 their degeneracy weights (15 per line), then one line
188 ``R1 R2 R3 m n Re Im`` per matrix element
189 :math:`H_{mn}(\mathbf{R}) = \langle m\mathbf{0}|H|n\mathbf{R}\rangle`
190 (m, n counted from 1).
192 Parameters
193 ----------
194 path : String or path
195 The ``_hr.dat`` file.
197 Returns
198 -------
199 num_wann : Integer
200 Number of Wannier functions.
201 R : Integer ndarray, shape (nrpts, 3)
202 Lattice vectors (in
203 units of the primitive vectors).
204 ndegen : Integer ndarray, shape (nrpts,)
205 Degeneracy weights.
206 ham : Complex ndarray, shape (nrpts, num_wann, num_wann)
207 :math:`H_{mn}(\mathbf{R})` as written in the file (not divided by
208 the weights).
209 '''
210 error_handling.file_path(path, 'path')
211 with open(path) as f:
212 lines = f.read().splitlines()[1:] # the first line is a free comment
213 tokens = ' '.join(lines).split()
214 error_handling.hr_header(tokens)
215 num_wann, nrpts = int(tokens[0]), int(tokens[1])
216 ndegen = np.array(tokens[2:2 + nrpts], dtype='i8')
217 body = tokens[2 + nrpts:]
218 error_handling.hr_body(len(body), nrpts, num_wann)
219 rows = np.array(body, dtype='f8').reshape(nrpts * num_wann ** 2, 7)
220 R = rows[::num_wann ** 2, :3].astype('i8')
221 block = np.repeat(np.arange(nrpts), num_wann ** 2)
222 error_handling.hr_blocks(np.array_equal(rows[:, :3].astype('i8'), R[block]))
223 m, n = rows[:, 3].astype('i8') - 1, rows[:, 4].astype('i8') - 1
224 ham = np.zeros((nrpts, num_wann, num_wann), 'c16')
225 ham[block, m, n] = rows[:, 5] + 1j * rows[:, 6]
226 return num_wann, R, ndegen, ham
229def read_win_cell(path: str | os.PathLike) -> NDArray[np.float64]:
230 r'''
231 Read the lattice vectors of a Wannier90 ``seedname.win`` file (block
232 ``unit_cell_cart``, in Angstrom, or in bohr if its first line says
233 ``bohr``). Comments (``!``, ``#``) are ignored.
235 Parameters
236 ----------
237 path : String or path
238 The ``.win`` file.
240 Returns
241 -------
242 cell : Real ndarray, shape (3, 3)
243 The rows are
244 :math:`\mathbf{a}_1, \mathbf{a}_2, \mathbf{a}_3`, in Angstrom.
245 '''
246 error_handling.file_path(path, 'path')
247 with open(path) as f:
248 lines = [line.split('!')[0].split('#')[0].strip() for line in f]
249 lines = [line for line in lines if line]
250 low = [line.lower().replace(' ', '') for line in lines]
251 error_handling.win_block('beginunit_cell_cart' in low and 'endunit_cell_cart' in low)
252 block = lines[low.index('beginunit_cell_cart') + 1:low.index('endunit_cell_cart')]
253 scale = 1.
254 if block and block[0].lower() in ('bohr', 'ang'):
255 scale = BOHR if block[0].lower() == 'bohr' else 1.
256 block = block[1:]
257 cell = [line.split() for line in block]
258 error_handling.win_cell(cell)
259 return scale * np.array(cell, dtype='f8')
262def read_centres(path: str | os.PathLike, num_wann: int) -> NDArray[np.float64]:
263 '''
264 Read the Wannier centres of a Wannier90 ``seedname_centres.xyz`` file
265 (the first *num_wann* atoms labelled ``X``, in Angstrom).
267 Parameters
268 ----------
269 path : String or path
270 The ``_centres.xyz`` file.
271 num_wann : Positive integer
272 Number of Wannier functions.
274 Returns
275 -------
276 centres : Real ndarray, shape (num_wann, 3)
277 '''
278 error_handling.file_path(path, 'path')
279 error_handling.positive_int(num_wann, 'num_wann')
280 with open(path) as f:
281 lines = f.read().splitlines()[2:] # atom count and comment
282 xs = [line.split()[1:4] for line in lines if line.split() and line.split()[0] == 'X']
283 error_handling.centres_count(len(xs), num_wann)
284 return np.array(xs[:num_wann], dtype='f8')
287def read_wannier90(
288 hr: str | os.PathLike,
289 prim_vec: list[tuple[float, ...]] | None = None,
290 win: str | os.PathLike | None = None,
291 centres: str | os.PathLike | None = None,
292 positions: ArrayLike | None = None,
293 tags: list[str] | None = None,
294 dim: int = 3,
295 cutoff: float = 0.,
296 tol: float = 1e-6,
297) -> KSpace:
298 r'''
299 Import a tight-binding model from Wannier90's ``seedname_hr.dat``,
300 the Hamiltonian in a basis of (maximally localized) Wannier functions:
302 .. math::
304 H_{mn}(\mathbf{k}) = \sum_{\mathbf{R}} \frac{H_{mn}(\mathbf{R})}{N_{\mathbf{R}}}\,
305 e^{i\mathbf{k}\cdot\mathbf{R}}\, ,
307 :math:`N_{\mathbf{R}}` being the degeneracy weight of :math:`\mathbf{R}`
308 (lattice vectors on the boundary of the Wigner-Seitz supercell appear
309 several times). The Wannier functions become the orbitals of a
310 **KSpace**, in the file's order: :math:`H_{mm}(\mathbf{0})` the onsite
311 energies, the other elements the hoppings
312 :math:`t_{mn}(\mathbf{R})` of *KSpace.set_hopping*. If the file is
313 Hermitian (within *tol*), one element of each conjugate pair is kept
314 and the model is Hermitian (with real onsite energies); otherwise every
315 element is kept, with ``hermitian=False``.
317 Parameters
318 ----------
319 hr : String or path
320 The ``seedname_hr.dat`` file.
321 prim_vec : List of *dim* tuples of *dim* real numbers
322 Default
323 value None. Primitive vectors, in the units of the positions
324 (Angstrom for Wannier90). If None, read from *win*.
325 win : String or path, default None
326 The ``seedname.win``
327 file, for its ``unit_cell_cart`` block (if *prim_vec* is None).
328 centres : String or path, default None
329 The
330 ``seedname_centres.xyz`` file: the Wannier centres become the
331 orbital positions (used by the Wilson-loop tools, the Peierls
332 phases, *bridges.finite_system*, ...).
333 positions : Array, shape (num_wann, dim) or (num_wann, 3), default None
334 Orbital positions, instead of *centres*. With
335 neither, every orbital sits at the origin of the cell.
336 tags : List of num_wann one-character strings
337 Default value
338 None (every orbital tagged ``'a'``). Sublattice tags.
339 dim : 2 or 3, default 3
340 A 2D model (a layer) keeps only
341 :math:`(R_1, R_2)`: every element with :math:`R_3\neq 0` must vanish
342 (within *tol*), and the lattice lies in the :math:`(x, y)` plane
343 (the z components of the vectors read from *win* must vanish; those
344 of the positions are dropped).
345 cutoff : Positive real number or zero, default 0
346 Drop the
347 hoppings with :math:`|t| \le` *cutoff* (the exact zeros by default).
348 tol : Positive real number, default 1e-6
349 Tolerance of the
350 Hermiticity and 2D checks (the file has 6 decimals).
352 Returns
353 -------
354 ks : **KSpace** instance, with num_wann orbitals
356 Examples
357 --------
358 Example usage::
360 ks = read_wannier90('graphene_hr.dat', win='graphene.win',
361 centres='graphene_centres.xyz', dim=2)
362 '''
363 error_handling.file_path(hr, 'hr')
364 error_handling.wannier_dim(dim)
365 error_handling.positive_real_zero(cutoff, 'cutoff')
366 error_handling.positive_real(tol, 'tol')
367 num_wann, R, ndegen, ham = read_hr(hr)
368 if prim_vec is None:
369 error_handling.win_given(win)
370 cell = read_win_cell(cast('str | os.PathLike', win))
371 error_handling.layer_cell(cell, dim)
372 prim_vec = [tuple(float(c) for c in a[:dim]) for a in cell[:dim]]
373 error_handling.prim_vec(prim_vec)
374 error_handling.wannier_prim_vec(prim_vec, dim)
375 if positions is None and centres is not None:
376 positions = read_centres(centres, num_wann)
377 if positions is None:
378 positions = np.zeros((num_wann, dim))
379 positions = np.asarray(positions, dtype='f8')
380 error_handling.wannier_positions(positions.shape, num_wann, dim)
381 tags = ['a'] * num_wann if tags is None else tags
382 error_handling.wannier_tags(tags, num_wann)
383 if dim == 2:
384 error_handling.layer_hoppings(float(np.abs(ham[R[:, 2] != 0]).max(initial=0.)), tol)
385 keep = R[:, 2] == 0
386 R, ndegen, ham = R[keep, :2], ndegen[keep], ham[keep]
387 return _hr_kspace(prim_vec, positions[:, :dim], tags, R, ham / ndegen[:, None, None], cutoff, tol)
390def _hr_kspace(
391 prim_vec: list[tuple[float, ...]], positions: NDArray[np.float64], tags: list[str],
392 R: NDArray[np.int_], ham: NDArray[np.complex128], cutoff: float, tol: float,
393) -> KSpace:
394 r'''
395 Private function. The **KSpace** of the matrices :math:`H(\mathbf{R})`
396 (*ham*, shape (nR, num_wann, num_wann), degeneracy weights already
397 divided out) at the lattice vectors *R* (integers, shape (nR, dim)),
398 with orbitals at *positions*, for *read_wannier90* and
399 *WannierFunctions.kspace*. It is Hermitian (with real onsite energies)
400 when every :math:`H(\mathbf{R})` has its partner
401 :math:`H(-\mathbf{R}) = H(\mathbf{R})^\dagger` within *tol*.
402 '''
403 num_wann = ham.shape[1]
404 unit_cell = [{'tag': tag, 'r0': tuple(float(c) for c in r)} for tag, r in zip(tags, positions)]
405 ks = KSpace(Lattice(unit_cell=unit_cell, prim_vec=list(prim_vec)))
406 # Hermitian: every H(R) has its partner H(-R) = H(R)^dagger
407 index = {tuple(r): p for p, r in enumerate(R.tolist())}
408 partner = [index.get(tuple(-c for c in r)) for r in R.tolist()]
409 hermitian = None not in partner and bool(
410 np.abs(ham - ham[np.array(partner)].conj().transpose(0, 2, 1)).max() <= tol)
411 list_hop = []
412 for r, h in zip(R.tolist(), ham):
413 nonzero = [c for c in r if c != 0]
414 for m in range(num_wann):
415 for n in range(num_wann):
416 t = complex(h[m, n])
417 if m == n and not nonzero:
418 ks.onsite[m] = t.real if hermitian else t
419 elif abs(t) > cutoff and (not hermitian or (nonzero and nonzero[0] > 0)
420 or (not nonzero and m < n)):
421 list_hop.append({'i': m, 'j': n, 'R': tuple(r), 't': t})
422 ks.set_hopping(list_hop, hermitian=hermitian)
423 return ks