Coverage for tbkit/io.py: 100%
154 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 06:13 +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 :param model: **Lattice**, **System** or **KSpace** instance.
80 :param path: String or path. File name; ``.npz`` is appended if missing.
82 :returns:
83 * **path** -- String. The file written.
85 Example usage::
87 save_model(gra, 'graphene.npz')
88 gra2 = load_model('graphene.npz')
89 '''
90 error_handling.saveable_model(model, Lattice, System, KSpace, OrbitalSystem)
91 error_handling.no_value_functions(model, 'save_model')
92 error_handling.file_path(path, 'path')
93 path = os.fspath(path)
94 if not path.endswith('.npz'):
95 path += '.npz'
96 lat = model.lat if isinstance(model, (System, KSpace)) else model
97 kind = 'System' if isinstance(model, System) else 'KSpace' if isinstance(model, KSpace) else 'Lattice'
98 # only what needs version 2 (a periodic System) is written as such, so
99 # that older tbkit versions keep reading everything else
100 data: dict[str, Any] = {'format': np.array(FORMAT), 'version': np.array(1), 'kind': np.array(kind),
101 'tbkit_version': np.array(tbkit.__version__),
102 'uc_tag': np.array([dic['tag'] for dic in lat.unit_cell], dtype='U1'),
103 'uc_r0': np.array([dic['r0'] for dic in lat.unit_cell], dtype='f8'),
104 'prim_vec': np.array(lat.prim_vec, dtype='f8'),
105 'coor': lat.coor,
106 'n_cells': np.array([lat.n1, lat.n2, lat.n3], dtype='i8')}
107 if isinstance(model, System):
108 data['onsite'] = np.asarray(model.onsite, dtype='c16')
109 data['hop'] = model.hop
110 if any(model.periodic):
111 data['periodic'] = np.array(model.periodic, dtype=bool)
112 data['version'] = np.array(VERSION)
113 elif isinstance(model, KSpace):
114 data.update({'spin': np.array(model.spin), 'onsite': model.onsite,
115 'onsite_offdiag': model._onsite_offdiag,
116 'nonreciprocal': np.array(model._nonreciprocal)})
117 data.update(_hop_arrays(model._hop, model.space_dim, 'hop'))
118 data.update(_hop_arrays(model._overlap_hop, model.space_dim, 'ovl'))
119 np.savez(path, **data)
120 return path
123def load_model(path: str | os.PathLike) -> Lattice | System | KSpace:
124 '''
125 Load a model saved by *save_model*.
127 :param path: String or path. File name.
129 :returns:
130 * **model** -- **Lattice**, **System** or **KSpace** instance, whose
131 Hamiltonians (*get_ham*) are exactly those of the saved model.
132 '''
133 error_handling.file_path(path, 'path')
134 with np.load(path, allow_pickle=False) as data:
135 error_handling.model_archive(data, FORMAT, VERSION)
136 unit_cell = [{'tag': str(tag), 'r0': tuple(float(c) for c in r0)}
137 for tag, r0 in zip(data['uc_tag'], data['uc_r0'])]
138 prim_vec = [tuple(float(c) for c in a) for a in data['prim_vec']]
139 lat = Lattice(unit_cell=unit_cell, prim_vec=prim_vec)
140 lat.coor = data['coor'].astype(lat.dtype)
141 lat.sites = len(lat.coor)
142 lat.n1, lat.n2, lat.n3 = (int(n) for n in data['n_cells'])
143 if lat.sites:
144 lat.tags = np.unique(np.concatenate([lat.tags, lat.coor['tag']]))
145 kind = str(data['kind'])
146 if kind == 'Lattice':
147 return lat
148 if kind == 'System':
149 periodic = tuple(bool(p) for p in data['periodic']) if 'periodic' in data.files else False
150 sys = System(lat, periodic=periodic)
151 sys.onsite = data['onsite'].copy()
152 sys.hop = data['hop'].copy()
153 return sys
154 ks = KSpace(lat, spin=bool(data['spin']))
155 ks.onsite = data['onsite'].copy()
156 ks._onsite_offdiag = data['onsite_offdiag'].copy()
157 ks._nonreciprocal = bool(data['nonreciprocal'])
158 ks._hop = _hop_list(data, 'hop')
159 ks._overlap_hop = _hop_list(data, 'ovl')
160 return ks
163#################################
164# WANNIER90
165#################################
167#: Bohr radius in Angstrom (CODATA 2018), for ``unit_cell_cart`` given in bohr.
168BOHR = 0.529177210903
171def read_hr(path: str | os.PathLike) -> tuple[int, NDArray[np.int64], NDArray[np.int64], NDArray[np.complex128]]:
172 r'''
173 Read a Wannier90 ``seedname_hr.dat`` file: a header line, the number of
174 Wannier functions, the number of lattice vectors :math:`\mathbf{R}`,
175 their degeneracy weights (15 per line), then one line
176 ``R1 R2 R3 m n Re Im`` per matrix element
177 :math:`H_{mn}(\mathbf{R}) = \langle m\mathbf{0}|H|n\mathbf{R}\rangle`
178 (m, n counted from 1).
180 :param path: String or path. The ``_hr.dat`` file.
182 :returns:
183 * **num_wann** -- Integer. Number of Wannier functions.
184 * **R** -- Integer ndarray, shape (nrpts, 3). Lattice vectors (in
185 units of the primitive vectors).
186 * **ndegen** -- Integer ndarray, shape (nrpts,). Degeneracy weights.
187 * **ham** -- Complex ndarray, shape (nrpts, num_wann, num_wann).
188 :math:`H_{mn}(\mathbf{R})` as written in the file (not divided by
189 the weights).
190 '''
191 error_handling.file_path(path, 'path')
192 with open(path) as f:
193 lines = f.read().splitlines()[1:] # the first line is a free comment
194 tokens = ' '.join(lines).split()
195 error_handling.hr_header(tokens)
196 num_wann, nrpts = int(tokens[0]), int(tokens[1])
197 ndegen = np.array(tokens[2:2 + nrpts], dtype='i8')
198 body = tokens[2 + nrpts:]
199 error_handling.hr_body(len(body), nrpts, num_wann)
200 rows = np.array(body, dtype='f8').reshape(nrpts * num_wann ** 2, 7)
201 R = rows[::num_wann ** 2, :3].astype('i8')
202 block = np.repeat(np.arange(nrpts), num_wann ** 2)
203 error_handling.hr_blocks(np.array_equal(rows[:, :3].astype('i8'), R[block]))
204 m, n = rows[:, 3].astype('i8') - 1, rows[:, 4].astype('i8') - 1
205 ham = np.zeros((nrpts, num_wann, num_wann), 'c16')
206 ham[block, m, n] = rows[:, 5] + 1j * rows[:, 6]
207 return num_wann, R, ndegen, ham
210def read_win_cell(path: str | os.PathLike) -> NDArray[np.float64]:
211 r'''
212 Read the lattice vectors of a Wannier90 ``seedname.win`` file (block
213 ``unit_cell_cart``, in Angstrom, or in bohr if its first line says
214 ``bohr``). Comments (``!``, ``#``) are ignored.
216 :param path: String or path. The ``.win`` file.
218 :returns:
219 * **cell** -- Real ndarray, shape (3, 3). The rows are
220 :math:`\mathbf{a}_1, \mathbf{a}_2, \mathbf{a}_3`, in Angstrom.
221 '''
222 error_handling.file_path(path, 'path')
223 with open(path) as f:
224 lines = [line.split('!')[0].split('#')[0].strip() for line in f]
225 lines = [line for line in lines if line]
226 low = [line.lower().replace(' ', '') for line in lines]
227 error_handling.win_block('beginunit_cell_cart' in low and 'endunit_cell_cart' in low)
228 block = lines[low.index('beginunit_cell_cart') + 1:low.index('endunit_cell_cart')]
229 scale = 1.
230 if block and block[0].lower() in ('bohr', 'ang'):
231 scale = BOHR if block[0].lower() == 'bohr' else 1.
232 block = block[1:]
233 cell = [line.split() for line in block]
234 error_handling.win_cell(cell)
235 return scale * np.array(cell, dtype='f8')
238def read_centres(path: str | os.PathLike, num_wann: int) -> NDArray[np.float64]:
239 '''
240 Read the Wannier centres of a Wannier90 ``seedname_centres.xyz`` file
241 (the first *num_wann* atoms labelled ``X``, in Angstrom).
243 :param path: String or path. The ``_centres.xyz`` file.
244 :param num_wann: Positive integer. Number of Wannier functions.
246 :returns:
247 * **centres** -- Real ndarray, shape (num_wann, 3).
248 '''
249 error_handling.file_path(path, 'path')
250 error_handling.positive_int(num_wann, 'num_wann')
251 with open(path) as f:
252 lines = f.read().splitlines()[2:] # atom count and comment
253 xs = [line.split()[1:4] for line in lines if line.split() and line.split()[0] == 'X']
254 error_handling.centres_count(len(xs), num_wann)
255 return np.array(xs[:num_wann], dtype='f8')
258def read_wannier90(
259 hr: str | os.PathLike,
260 prim_vec: list[tuple[float, ...]] | None = None,
261 win: str | os.PathLike | None = None,
262 centres: str | os.PathLike | None = None,
263 positions: ArrayLike | None = None,
264 tags: list[str] | None = None,
265 dim: int = 3,
266 cutoff: float = 0.,
267 tol: float = 1e-6,
268) -> KSpace:
269 r'''
270 Import a tight-binding model from Wannier90's ``seedname_hr.dat``,
271 the Hamiltonian in a basis of (maximally localized) Wannier functions:
273 .. math::
275 H_{mn}(\mathbf{k}) = \sum_{\mathbf{R}} \frac{H_{mn}(\mathbf{R})}{N_{\mathbf{R}}}\,
276 e^{i\mathbf{k}\cdot\mathbf{R}}\, ,
278 :math:`N_{\mathbf{R}}` being the degeneracy weight of :math:`\mathbf{R}`
279 (lattice vectors on the boundary of the Wigner-Seitz supercell appear
280 several times). The Wannier functions become the orbitals of a
281 **KSpace**, in the file's order: :math:`H_{mm}(\mathbf{0})` the onsite
282 energies, the other elements the hoppings
283 :math:`t_{mn}(\mathbf{R})` of *KSpace.set_hopping*. If the file is
284 Hermitian (within *tol*), one element of each conjugate pair is kept
285 and the model is Hermitian (with real onsite energies); otherwise every
286 element is kept, with ``hermitian=False``.
288 :param hr: String or path. The ``seedname_hr.dat`` file.
289 :param prim_vec: List of *dim* tuples of *dim* real numbers. Default
290 value None. Primitive vectors, in the units of the positions
291 (Angstrom for Wannier90). If None, read from *win*.
292 :param win: String or path. Default value None. The ``seedname.win``
293 file, for its ``unit_cell_cart`` block (if *prim_vec* is None).
294 :param centres: String or path. Default value None. The
295 ``seedname_centres.xyz`` file: the Wannier centres become the
296 orbital positions (used by the Wilson-loop tools, the Peierls
297 phases, *bridges.finite_system*, ...).
298 :param positions: Array, shape (num_wann, dim) or (num_wann, 3).
299 Default value None. Orbital positions, instead of *centres*. With
300 neither, every orbital sits at the origin of the cell.
301 :param tags: List of num_wann one-character strings. Default value
302 None (every orbital tagged ``'a'``). Sublattice tags.
303 :param dim: 2 or 3. Default value 3. A 2D model (a layer) keeps only
304 :math:`(R_1, R_2)`: every element with :math:`R_3\neq 0` must vanish
305 (within *tol*), and the lattice lies in the :math:`(x, y)` plane
306 (the z components of the vectors read from *win* must vanish; those
307 of the positions are dropped).
308 :param cutoff: Positive real number or zero. Default value 0. Drop the
309 hoppings with :math:`|t| \le` *cutoff* (the exact zeros by default).
310 :param tol: Positive real number. Default value 1e-6. Tolerance of the
311 Hermiticity and 2D checks (the file has 6 decimals).
313 :returns:
314 * **ks** -- **KSpace** instance, with num_wann orbitals.
316 Example usage::
318 ks = read_wannier90('graphene_hr.dat', win='graphene.win',
319 centres='graphene_centres.xyz', dim=2)
320 '''
321 error_handling.file_path(hr, 'hr')
322 error_handling.wannier_dim(dim)
323 error_handling.positive_real_zero(cutoff, 'cutoff')
324 error_handling.positive_real(tol, 'tol')
325 num_wann, R, ndegen, ham = read_hr(hr)
326 if prim_vec is None:
327 error_handling.win_given(win)
328 cell = read_win_cell(cast('str | os.PathLike', win))
329 error_handling.layer_cell(cell, dim)
330 prim_vec = [tuple(float(c) for c in a[:dim]) for a in cell[:dim]]
331 error_handling.prim_vec(prim_vec)
332 error_handling.wannier_prim_vec(prim_vec, dim)
333 if positions is None and centres is not None:
334 positions = read_centres(centres, num_wann)
335 if positions is None:
336 positions = np.zeros((num_wann, dim))
337 positions = np.asarray(positions, dtype='f8')
338 error_handling.wannier_positions(positions.shape, num_wann, dim)
339 tags = ['a'] * num_wann if tags is None else tags
340 error_handling.wannier_tags(tags, num_wann)
341 if dim == 2:
342 error_handling.layer_hoppings(float(np.abs(ham[R[:, 2] != 0]).max(initial=0.)), tol)
343 keep = R[:, 2] == 0
344 R, ndegen, ham = R[keep, :2], ndegen[keep], ham[keep]
345 return _hr_kspace(prim_vec, positions[:, :dim], tags, R, ham / ndegen[:, None, None], cutoff, tol)
348def _hr_kspace(
349 prim_vec: list[tuple[float, ...]], positions: NDArray[np.float64], tags: list[str],
350 R: NDArray[np.int_], ham: NDArray[np.complex128], cutoff: float, tol: float,
351) -> KSpace:
352 r'''
353 Private function. The **KSpace** of the matrices :math:`H(\mathbf{R})`
354 (*ham*, shape (nR, num_wann, num_wann), degeneracy weights already
355 divided out) at the lattice vectors *R* (integers, shape (nR, dim)),
356 with orbitals at *positions*, for *read_wannier90* and
357 *WannierFunctions.kspace*. It is Hermitian (with real onsite energies)
358 when every :math:`H(\mathbf{R})` has its partner
359 :math:`H(-\mathbf{R}) = H(\mathbf{R})^\dagger` within *tol*.
360 '''
361 num_wann = ham.shape[1]
362 unit_cell = [{'tag': tag, 'r0': tuple(float(c) for c in r)} for tag, r in zip(tags, positions)]
363 ks = KSpace(Lattice(unit_cell=unit_cell, prim_vec=list(prim_vec)))
364 # Hermitian: every H(R) has its partner H(-R) = H(R)^dagger
365 index = {tuple(r): p for p, r in enumerate(R.tolist())}
366 partner = [index.get(tuple(-c for c in r)) for r in R.tolist()]
367 hermitian = None not in partner and bool(
368 np.abs(ham - ham[np.array(partner)].conj().transpose(0, 2, 1)).max() <= tol)
369 list_hop = []
370 for r, h in zip(R.tolist(), ham):
371 nonzero = [c for c in r if c != 0]
372 for m in range(num_wann):
373 for n in range(num_wann):
374 t = complex(h[m, n])
375 if m == n and not nonzero:
376 ks.onsite[m] = t.real if hermitian else t
377 elif abs(t) > cutoff and (not hermitian or (nonzero and nonzero[0] > 0)
378 or (not nonzero and m < n)):
379 list_hop.append({'i': m, 'j': n, 'R': tuple(r), 't': t})
380 ks.set_hopping(list_hop, hermitian=hermitian)
381 return ks