Coverage for tbkit/wannier.py: 100%
241 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-10-30 11:04 +0000
1r"""
2Wannier functions of a group of bands of a :class:`tbkit.kspace.KSpace`
3model, built in two steps (three for entangled bands):
5* **Disentanglement** (when there are more bands than Wannier functions):
6 at each k, the :math:`n_W`-dimensional subspace of the states in an outer
7 energy window that is smoothest across the mesh, i.e. that minimizes
8 :math:`\Omega_I`, keeping the states of an optional frozen window exactly
9 (Souza, Marzari and Vanderbilt, Phys. Rev. B 65, 035109 (2001)).
10* **Projection** onto trial orbitals :math:`|g_n\rangle` localized in the
11 home cell: :math:`A_{mn}(\mathbf{k}) = \langle u_{m\mathbf{k}}|g_n\rangle`,
12 made unitary by Lowdin orthogonalization,
13 :math:`U(\mathbf{k}) = A(A^\dagger A)^{-1/2}` (Marzari, Mostofi, Yates,
14 Souza and Vanderbilt, Rev. Mod. Phys. 84, 1419 (2012), Sec. III.A).
15* **Maximal localization**: conjugate-gradient descent of the spread
16 :math:`\Omega = \sum_n \langle r^2\rangle_n - |\bar{\mathbf{r}}_n|^2`
17 over the gauge :math:`U(\mathbf{k})`, from the overlaps
18 :math:`M^{(\mathbf{k},\mathbf{b})}_{mn} = \langle u_{m\mathbf{k}}|u_{n\mathbf{k}+\mathbf{b}}\rangle`
19 between neighbouring points of the k-mesh (Marzari and Vanderbilt, Phys.
20 Rev. B 56, 12847 (1997)).
22The Wannier functions are the Fourier transforms of the rotated Bloch
23states, on the :math:`N_1\times N_2\times\dots` cells of the k-mesh (the
24supercell on which they are periodic). Their Hamiltonian
25:math:`H_{mn}(\mathbf{R}) = \langle W_{m\mathbf{0}}|H|W_{n\mathbf{R}}\rangle`
26is a tight-binding model of its own (*WannierFunctions.kspace*), which
27interpolates the bands between the mesh points (Wannier interpolation).
29A band group with a nonzero Chern number has no exponentially localized
30Wannier functions (Thouless, J. Phys. C 17, L325 (1984); Brouder et al.,
31Phys. Rev. Lett. 98, 046402 (2007)): any projection vanishes somewhere in
32the Brillouin zone, and the spread grows with the mesh instead of
33converging.
34"""
35from __future__ import annotations
37import numpy as np
38from numpy.typing import ArrayLike, NDArray
40import tbkit.error_handling as error_handling
41from tbkit.bridges import _finite_sites, finite_system
42from tbkit.io import _hr_kspace
43from tbkit.kspace import KSpace
46class WannierFunctions():
47 r'''
48 Wannier functions of a group of bands, see *wannierize*.
50 Attributes
51 ----------
52 functions : Complex ndarray, shape (n_wann, N*norb)
53 Amplitudes
54 of each Wannier function on the orbitals of the
55 :math:`N = N_1N_2\dots` cells of the k-mesh, indexed as the rows of
56 *KSpace.finite_ham(nk)* (and the sites of *system*). Each one is a
57 lattice translate chosen to sit in the middle of the sample.
58 positions : Real ndarray, shape (N*norb, space_dim)
59 Position of each row.
60 centers : Real ndarray, shape (n_wann, space_dim)
61 Centres
62 :math:`\bar{\mathbf{r}}_n` of *functions* (Cartesian).
63 spreads : Real ndarray, shape (n_wann,)
64 Spreads
65 :math:`\langle r^2\rangle_n - |\bar{\mathbf{r}}_n|^2` (squared
66 length), along the primitive vectors only for a chain or a 2D
67 lattice in 3D.
68 omega_i : Real number
69 Gauge-invariant part :math:`\Omega_I` of the total spread.
70 omega_d : Real number
71 Diagonal part :math:`\tilde\Omega_D`.
72 omega_od : Real number
73 Off-diagonal part :math:`\tilde\Omega_{OD}`.
74 gauge : Complex ndarray, shape (N, len(bands), n_wann)
75 :math:`U(\mathbf{k})` over the k-mesh (flattened as *KSpace.mesh_grid*),
76 acting on the eigenvectors of *bands*: unitary for an isolated
77 group, otherwise the disentangled subspace (orthonormal columns,
78 zero on the states outside the window) times the unitary gauge.
79 history : Real ndarray
80 Total spread after the projection and after each step.
81 converged : Boolean
82 Whether the spread changed by less than *tol*.
83 min_singular_value : Real number
84 Smallest singular value of the
85 projection :math:`A(\mathbf{k})` over the mesh: close to zero when the
86 trial orbitals miss the bands somewhere (as they must for a Chern band).
87 dis_history : Real ndarray
88 :math:`\Omega_I` after the projection
89 and after each disentanglement step (empty for an isolated group).
90 dis_converged : Boolean
91 Whether :math:`\Omega_I` changed by
92 less than *dis_tol* (True for an isolated group).
93 ham_k : Complex ndarray, shape (N, n_wann, n_wann)
94 Hamiltonian
95 :math:`U^\dagger(\mathbf{k})\,\mathrm{diag}(E)\,U(\mathbf{k})`
96 in the gauge of the Wannier functions, over the k-mesh.
97 '''
99 def __init__(self, ks, n_cells, functions, positions, centers, spreads, omega_i,
100 omega_d, omega_od, gauge, history, converged, min_singular_value,
101 dis_history=None, dis_converged=True, ham_k=None, shift=None) -> None:
102 self._ks = ks
103 self.n_cells = n_cells
104 self.functions = functions
105 self.positions = positions
106 self.centers = centers
107 self.spreads = spreads
108 self.omega_i = omega_i
109 self.omega_d = omega_d
110 self.omega_od = omega_od
111 self.gauge = gauge
112 self.history = history
113 self.converged = converged
114 self.min_singular_value = min_singular_value
115 self.dis_history = np.array([]) if dis_history is None else dis_history
116 self.dis_converged = dis_converged
117 self.ham_k = ham_k
118 self._shift = shift # lattice vectors (integers) moving each function to the middle cell
119 self._system = None
121 @property
122 def omega(self) -> float:
123 r'''
124 Total spread :math:`\Omega = \Omega_I + \tilde\Omega_D + \tilde\Omega_{OD}`,
125 the sum of *spreads*.
126 '''
127 return float(np.sum(self.spreads))
129 @property
130 def system(self):
131 '''
132 The finite *System* (open boundaries) whose sites carry *functions*:
133 *bridges.finite_system(ks, nk)*, built on first access. Spinless
134 models only (for a spinful one, use *positions*).
135 '''
136 if self._system is None:
137 self._system = finite_system(self._ks, self.n_cells)
138 return self._system
140 def kspace(self, cutoff: float = 1e-10) -> KSpace:
141 r'''
142 The tight-binding model of the Wannier functions (Wannier
143 interpolation): one orbital per function, at its centre (folded
144 into the home cell), and the matrix elements
146 .. math::
148 H_{mn}(\mathbf{R}) = \langle W_{m\mathbf{0}}|H|W_{n\mathbf{R}}\rangle
149 = \frac{1}{N}\sum_{\mathbf{k}} e^{-i\mathbf{k}\cdot\mathbf{R}}\,
150 [U^\dagger(\mathbf{k})\,\mathrm{diag}(E)\,U(\mathbf{k})]_{mn}
152 as onsite energies and hoppings. Each :math:`\mathbf{R}` is the
153 image on the :math:`N_1\times N_2\times\dots` supercell closest to
154 the bond, :math:`|\mathbf{R} + \boldsymbol\tau_n - \boldsymbol\tau_m|`
155 minimal (equidistant images share the element, as
156 Wannier90's ``use_ws_distance``). Its bands equal the Wannierized
157 ones on the k-mesh: all of them for an isolated group, those in the
158 frozen window after disentanglement. Between the mesh points they
159 converge exponentially with the mesh, as the functions decay. The
160 model is the one *io.read_wannier90* would read from Wannier90's
161 ``seedname_hr.dat``.
163 Parameters
164 ----------
165 cutoff : Positive real number or zero, default 1e-10
166 Drop the hoppings with :math:`|t| \le` *cutoff*.
168 Returns
169 -------
170 ks : **KSpace** instance, Hermitian, with n_wann orbitals
172 Examples
173 --------
174 Example usage::
176 wf = wannierize(ks, list(range(ks.norb)), trial=[0, 1], frozen=(-3., 0.))
177 wk = wf.kspace()
178 dist, en = wk.k_path(high_symmetry_path(ks.lat)[0], 50)
179 '''
180 error_handling.positive_real_zero(cutoff, 'cutoff')
181 ks, nk = self._ks, tuple(self.n_cells)
182 dim, n_k, n_w = ks.dim, int(np.prod(nk)), len(self.spreads)
183 prim = np.array(ks.lat.prim_vec, dtype='f8')
184 # orbital m is W_m(T_m), the translate of the Fourier transform W_m(0)
185 # (the functions before their move to the middle cell) in the home cell
186 centre = self.centers - self._shift @ prim
187 cell = np.floor(centre @ np.array(ks.rec_vec, dtype='f8').T / (2 * np.pi) + 1e-8).astype(int)
188 tau = centre - cell @ prim
189 # H'_mn(R) = H_mn(R + T_n - T_m): a phase on the mesh, then the Fourier transform
190 fracs = np.array(np.meshgrid(*[np.arange(n) / n for n in nk], indexing='ij')).reshape(dim, -1).T
191 dt = cell[None, :, :] - cell[:, None, :] # T_n - T_m, shape (n_w, n_w, dim)
192 ham = self.ham_k * np.exp(-2j * np.pi * np.einsum('ka,mna->kmn', fracs, dt))
193 ham_r = np.fft.fftn(ham.reshape(nk + (n_w, n_w)), axes=tuple(range(dim))) / n_k
194 # the image R = j + N s closest to each bond, j centred on 0
195 j = np.array(np.meshgrid(*[(np.arange(n) + n // 2) % n - n // 2 for n in nk],
196 indexing='ij')).reshape(dim, -1).T
197 s = np.array(np.meshgrid(*[np.arange(-1, 2)] * dim, indexing='ij')).reshape(dim, -1).T
198 R = (j[:, None, :] + s[None, :, :] * np.array(nk)).reshape(-1, dim) # (n_k * 3^dim, dim)
199 bond = (R @ prim)[:, None, None, :] + tau[None, None, :, :] - tau[None, :, None, :]
200 dist = np.linalg.norm(bond, axis=3).reshape(n_k, len(s), n_w, n_w)
201 closest = dist <= dist.min(axis=1, keepdims=True) + 1e-6 * max(1., float(np.abs(prim).max()))
202 weight = closest / np.sum(closest, axis=1, keepdims=True)
203 ham_r = (ham_r.reshape(n_k, 1, n_w, n_w) * weight).reshape(-1, n_w, n_w)
204 keep = np.any(weight.reshape(-1, n_w, n_w) > 0., axis=(1, 2))
205 return _hr_kspace(ks.lat.prim_vec, tau, ['a'] * n_w, R[keep], ham_r[keep], cutoff, 1e-8)
208def _shells(rec: NDArray[np.float64], nk: tuple[int, ...]) -> tuple[NDArray, NDArray, NDArray]:
209 r'''
210 Private function. Neighbours :math:`\mathbf{b}` of a k-point on the mesh
211 and their weights :math:`w_b`, so that
212 :math:`\sum_b w_b b_\alpha b_\beta = \delta_{\alpha\beta}` (the
213 finite-difference condition of Marzari and Vanderbilt 1997, Eq. B1).
214 Shells of equal :math:`|\mathbf{b}|` are added by increasing length,
215 skipping those that add no independent direction, until the condition
216 is solvable (as Wannier90 does).
218 Returns
219 -------
220 steps : integers, shape (nb, dim)
221 b : shape (nb, dim)
222 w : shape (nb,)
223 '''
224 dim = len(nk)
225 steps = np.array(np.meshgrid(*[np.arange(-1, 2)] * dim, indexing='ij')).reshape(dim, -1).T
226 steps = steps[np.any(steps != 0, axis=1)]
227 b = (steps / np.array(nk)[None, :]) @ rec
228 length = np.linalg.norm(b, axis=1)
229 radii = np.unique(np.round(length / length.min(), 8))
230 upper = np.triu_indices(dim)
231 target = np.eye(dim)[upper]
232 chosen: list = []
233 columns: list = []
234 for r in radii:
235 shell = np.abs(length / length.min() - r) < 1e-6
236 column = np.einsum('ba,bc->ac', b[shell], b[shell])[upper]
237 trial = np.array(columns + [column]).T
238 if np.linalg.matrix_rank(trial, tol=1e-10) < len(columns) + 1:
239 continue
240 chosen.append(shell)
241 columns.append(column)
242 w, *_ = np.linalg.lstsq(trial, target, rcond=None)
243 # always reached: the steps e_i and e_i + e_j give dim(dim+1)/2
244 # independent columns, and the system is then square and invertible
245 if np.allclose(trial @ w, target, atol=1e-8):
246 break
247 keep = np.concatenate([np.flatnonzero(s) for s in chosen])
248 weights = np.concatenate([np.full(np.sum(s), ws) for s, ws in zip(chosen, w)])
249 return steps[keep], b[keep], weights
252def _spread(m, b, w, n_k):
253 r'''
254 Private function. Centres (shape (n_wann, dim), k coordinates), spreads
255 and the parts :math:`\Omega_I, \tilde\Omega_D, \tilde\Omega_{OD}` from
256 the overlaps *m* (shape (nk, nb, n_wann, n_wann)).
257 '''
258 diag = np.diagonal(m, axis1=2, axis2=3) # (nk, nb, n_wann)
259 phase = np.angle(diag)
260 centers = -np.einsum('b,ba,kbn->na', w, b, phase) / n_k
261 r2 = np.einsum('b,kbn->n', w, 1. - np.abs(diag) ** 2 + phase ** 2) / n_k
262 spreads = r2 - np.sum(centers ** 2, axis=1)
263 total = np.einsum('b,kbmn->', w, np.abs(m) ** 2) / n_k
264 on_diag = np.einsum('b,kbn->', w, np.abs(diag) ** 2) / n_k
265 omega_i = float(np.sum(w) * m.shape[2] - total)
266 omega_od = float(total - on_diag)
267 omega_d = float(np.einsum('b,kbn->', w, (phase + (b @ centers.T)[None]) ** 2) / n_k)
268 return centers, spreads, omega_i, omega_d, omega_od
271def _gradient(m, b, w, centers):
272 r'''
273 Private function. Gradient :math:`G(\mathbf{k})` of the spread for a
274 gauge change :math:`U \to U e^{dW}` (Marzari and Vanderbilt 1997,
275 Eq. 52), :math:`G = 4\sum_b w_b(\mathcal{A}[R] - \mathcal{S}[T])`:
276 :math:`\Omega` decreases along :math:`dW = \epsilon G`, :math:`\epsilon > 0`,
277 with :math:`R_{mn} = M_{mn}M_{nn}^*`, :math:`T_{mn} = (M_{mn}/M_{nn})q_n`,
278 :math:`q_n = \mathrm{Im}\ln M_{nn} + \mathbf{b}\cdot\bar{\mathbf{r}}_n`,
279 :math:`\mathcal{A}[B] = (B - B^\dagger)/2`, :math:`\mathcal{S}[B] = (B + B^\dagger)/2i`.
280 Shape (nk, n_wann, n_wann), anti-Hermitian.
281 '''
282 diag = np.diagonal(m, axis1=2, axis2=3)
283 r = m * diag.conj()[:, :, None, :]
284 q = np.angle(diag) + (b @ centers.T)[None]
285 t = m / diag[:, :, None, :] * q[:, :, None, :]
286 a_r = (r - r.conj().transpose(0, 1, 3, 2)) / 2
287 s_t = (t + t.conj().transpose(0, 1, 3, 2)) / 2j
288 return 4 * np.einsum('b,kbmn->kmn', w, a_r - s_t)
291def _unitary_step(g: NDArray[np.complex128], eps: float) -> NDArray[np.complex128]:
292 '''
293 Private function. exp(eps G) for anti-Hermitian matrices G (stacked).
294 '''
295 h = 1j * g # Hermitian
296 lam, vec = np.linalg.eigh((h + h.conj().transpose(0, 2, 1)) / 2)
297 return (vec * np.exp(-1j * eps * lam)[:, None, :]) @ vec.conj().transpose(0, 2, 1)
300def _disentangle(m0, neighbour, w, inside, frozen, a, n_w, max_iter, tol, mixing):
301 r'''
302 Private function. Souza-Marzari-Vanderbilt subspace selection. At each
303 k, the :math:`n_W` states :math:`U(\mathbf{k})` (columns, in the basis
304 of the candidate bands) are the frozen ones plus the leading
305 eigenvectors of
307 .. math::
309 Z(\mathbf{k}) = \sum_b w_b\, M^{(0)}_{\mathbf{k},\mathbf{b}}\,
310 P(\mathbf{k}+\mathbf{b})\, M^{(0)\dagger}_{\mathbf{k},\mathbf{b}}
312 restricted to the free states of the window (*inside* and not *frozen*),
313 :math:`P = UU^\dagger` from the previous step, with linear mixing of
314 :math:`Z` (Eqs. 21-25 of SMV 2001). Each step lowers
315 :math:`\Omega_I = \sum_b w_b(n_W - \langle\|U_\mathbf{k}^\dagger M^{(0)}U_{\mathbf{k}+\mathbf{b}}\|^2\rangle_k)`.
316 The first subspace is the projection of the trial orbitals *a*
317 (overlaps :math:`\langle u_m|g_n\rangle`, shape (nk, N, n_w)).
319 Returns
320 -------
321 U : shape (nk, N, n_w)
322 history
323 The :math:`\Omega_I` history.
324 converged
325 '''
326 n_k, n = inside.shape
327 free = inside & ~frozen
328 # the eigenvalues of Z lie in [0, sum w]: frozen states above, excluded below
329 big = 2 * np.sum(w) + 1
330 level = np.where(frozen, big, np.where(inside, 0., -big))
331 diag = np.arange(n)
333 def select(z):
334 z = z * (free[:, :, None] & free[:, None, :])
335 z[:, diag, diag] += level
336 return np.linalg.eigh(z)[1][:, :, -n_w:]
338 def z_matrix(u):
339 mu = m0 @ u[neighbour] # (nk, nb, N, n_w)
340 return np.einsum('b,kbmi,kbni->kmn', w, mu, mu.conj())
342 u = select(a @ a.conj().transpose(0, 2, 1))
343 z = z_matrix(u)
344 history = [float(n_w * np.sum(w) - np.einsum('kmi,kmn,kni->', u.conj(), z, u).real / n_k)]
345 converged = max_iter == 0
346 z_in = z
347 for _ in range(max_iter):
348 u = select(z_in)
349 z = z_matrix(u)
350 history.append(float(n_w * np.sum(w) - np.einsum('kmi,kmn,kni->', u.conj(), z, u).real / n_k))
351 if abs(history[-2] - history[-1]) < tol:
352 converged = True
353 break
354 z_in = mixing * z + (1 - mixing) * z_in
355 return u, np.array(history), converged
358def wannierize(
359 ks: KSpace, bands: int | list[int], trial: list[int] | ArrayLike,
360 nk: int | tuple[int, ...] = 12, max_iter: int = 500, tol: float = 1e-10,
361 step: float = 0.5, window: tuple[float, float] | None = None,
362 frozen: tuple[float, float] | None = None, dis_iter: int = 2000, dis_tol: float = 1e-10,
363 mixing: float = 0.5,
364) -> WannierFunctions:
365 r'''
366 Build the maximally localized Wannier functions of a group of bands:
367 project the Bloch states onto trial orbitals (Lowdin orthogonalized),
368 then minimize the spread
370 .. math::
372 \Omega = \sum_n \left[\langle r^2\rangle_n - |\bar{\mathbf{r}}_n|^2\right]
373 = \Omega_I + \tilde\Omega_D + \tilde\Omega_{OD}
375 by conjugate gradients over the gauge :math:`U(\mathbf{k})` (Marzari
376 and Vanderbilt, Phys. Rev. B 56, 12847 (1997); as in Wannier90, with a
377 parabolic line search). The finite-difference
378 formulas use the overlaps of the Bloch states at neighbouring mesh
379 points, with the orbital positions included (as in the Wilson loops of
380 *KSpace.wannier_centers*, so the centres agree with them). Only
381 :math:`\tilde\Omega_D + \tilde\Omega_{OD}` depends on the gauge;
382 :math:`\Omega_I` is fixed by the band group.
384 In 1D, the minimum has :math:`\tilde\Omega_D = \tilde\Omega_{OD} = 0`
385 for one band, and the centre is the Berry phase over :math:`2\pi`. For
386 a Chern band, every projection fails somewhere (*min_singular_value*
387 tends to zero), and the spread diverges as the mesh is refined. The
388 descent finds the local minimum nearest the projected gauge: trial
389 orbitals that miss the bands at isolated k-points (a small
390 *min_singular_value*) leave phase vortices that it does not remove,
391 even in a trivial band, so choose orbitals the bands actually live on.
393 **Entangled bands.** With fewer trial orbitals :math:`n_W` than
394 *bands* (or with a *window*), the bands need not be isolated: the
395 Wannier functions span, at each k, the :math:`n_W`-dimensional subspace
396 of the states of *bands* inside the outer *window* that minimizes
397 :math:`\Omega_I` (Souza, Marzari and Vanderbilt, Phys. Rev. B 65,
398 035109 (2001)), found by iterating from the projection of the trial
399 orbitals, with linear *mixing*. The states inside the *frozen* window
400 are kept exactly, so the Wannier bands (*WannierFunctions.kspace*)
401 reproduce them. The maximal localization then runs within that
402 subspace. Pass every band (``list(range(ks.norb))``) and select the
403 states with the windows.
405 Parameters
406 ----------
407 ks : **KSpace** instance, Hermitian
408 bands : Band index, or list of band indices
409 With one trial
410 orbital per band and no window, the group, which must be separated
411 in energy from the other bands over the whole mesh; otherwise the
412 candidate bands of the disentanglement.
413 trial
414 Trial orbitals, one per Wannier function (at most one per
415 band): a list of orbital indices (each a single orbital of the home
416 cell), or a complex array of shape (norb, n_wann) whose columns are
417 combinations of the orbitals of the home cell (e.g. a bond orbital).
418 nk : Positive integer, or tuple of *dim* integers
419 Default value
420 12. k-mesh, at least 3 along each direction; the Wannier functions
421 live on the same number of cells.
422 max_iter : Positive integer or zero, default 500
423 Maximal
424 number of descent steps (0 keeps the projected gauge).
425 tol : Positive real number, default 1e-10
426 Convergence
427 threshold on the change of :math:`\Omega` between steps.
428 step : Positive real number, default 0.5
429 Trial step
430 :math:`\alpha` of the line search, :math:`dW = \alpha D/4\sum_b w_b`
431 along the search direction :math:`D`; it is halved whenever the
432 spread would grow.
433 window : Tuple of two real numbers
434 Default value None (every
435 state of *bands*). Outer energy window :math:`(E_{min}, E_{max})`:
436 the states of *bands* inside it are the candidates, at least
437 n_wann at every k-point.
438 frozen : Tuple of two real numbers
439 Default value None (no
440 frozen states). Inner energy window, inside *window*: its states,
441 at most n_wann at every k-point, are kept in the subspace.
442 dis_iter : Positive integer or zero, default 2000
443 Maximal
444 number of disentanglement steps (0 keeps the projected subspace).
445 dis_tol : Positive real number, default 1e-10
446 Convergence
447 threshold on the change of :math:`\Omega_I` between steps.
448 mixing : Real in (0, 1], default 0.5
449 Fraction of the new
450 :math:`Z(\mathbf{k})` mixed in at each disentanglement step.
452 Returns
453 -------
454 wf : **WannierFunctions** instance
456 Examples
457 --------
458 Example usage::
460 wf = wannierize(ssh, 0, trial=[0], nk=20)
461 wf.centers, wf.spreads, wf.functions
462 # two Wannier functions out of four bands, the lowest two kept exactly
463 wf = wannierize(ks, [0, 1, 2, 3], trial=[0, 1], window=(-4., 3.), frozen=(-4., 0.))
464 wf.kspace() # the Wannier-interpolated model
465 '''
466 error_handling.kspace(ks, KSpace)
467 error_handling.hermitian_model(ks.is_hermitian())
468 if isinstance(bands, int):
469 bands = [bands]
470 error_handling.band_indices(bands, ks.norb)
471 error_handling.trial_orbitals(trial, ks.norb, len(bands))
472 error_handling.nk(nk, ks.dim)
473 nk = (nk,) * ks.dim if isinstance(nk, int) else nk
474 error_handling.nk_min(nk, 3)
475 error_handling.positive_int_zero(max_iter, 'max_iter')
476 error_handling.positive_real(tol, 'tol')
477 error_handling.positive_real(step, 'step')
478 if window is not None:
479 error_handling.energy_window(window, 'window')
480 if frozen is not None:
481 error_handling.energy_window(frozen, 'frozen')
482 if window is not None:
483 error_handling.frozen_window(frozen, window)
484 error_handling.positive_int_zero(dis_iter, 'dis_iter')
485 error_handling.positive_real(dis_tol, 'dis_tol')
486 error_handling.mixing(mixing)
487 if isinstance(trial, list) and all(isinstance(o, int) for o in trial):
488 g = np.eye(ks.norb, dtype='c16')[:, trial]
489 else:
490 g = np.asarray(trial, dtype='c16')
492 n_k, n_w = int(np.prod(nk)), g.shape[1]
493 entangled = n_w < len(bands) or window is not None or frozen is not None
494 _, kpts = ks.mesh_grid(nk)
495 en = ks._eigs(kpts)
496 if not entangled:
497 error_handling.isolated_bands(en, bands)
498 v = ks._subspaces(kpts, bands) # (n_k, norb, len(bands))
500 # overlaps of the Bloch states with their mesh neighbours
501 steps, b, w = _shells(ks.rec_vec_k, nk)
502 index = np.arange(n_k).reshape(nk)
503 neighbour = np.stack([np.roll(index, tuple(-st), axis=tuple(range(ks.dim))).ravel()
504 for st in steps], axis=1) # (n_k, nb)
505 tau = ks.orbital_positions()
506 phase = np.exp(-1j * tau @ (ks.k_basis @ b.T)) # (norb, nb)
507 m0 = np.einsum('kom,ob,kbon->kbmn', v.conj(), phase, v[neighbour])
509 # disentanglement: the n_w-dimensional subspace, then the same steps within it
510 energies = en[:, bands]
511 dis_history, dis_converged = np.array([]), True
512 u_dis = np.broadcast_to(np.eye(n_w, dtype='c16'), (n_k, n_w, n_w))
513 if entangled:
514 lo, hi = window if window is not None else (-np.inf, np.inf)
515 inside = (energies >= lo) & (energies <= hi)
516 fixed = np.zeros_like(inside) if frozen is None else \
517 inside & (energies >= frozen[0]) & (energies <= frozen[1])
518 error_handling.window_states(int(inside.sum(axis=1).min()), int(fixed.sum(axis=1).max()), n_w)
519 u_dis, dis_history, dis_converged = _disentangle(
520 m0, neighbour, w, inside, fixed, v.conj().transpose(0, 2, 1) @ g, n_w,
521 dis_iter, dis_tol, mixing)
522 v = v @ u_dis # (n_k, norb, n_w)
523 m0 = u_dis.conj().transpose(0, 2, 1)[:, None] @ m0 @ u_dis[neighbour]
525 # projection, made unitary: A = W s V^dag -> U = W V^dag
526 a = v.conj().transpose(0, 2, 1) @ g
527 w_svd, s, vh = np.linalg.svd(a)
528 error_handling.projection(s.min())
529 u = w_svd @ vh
531 def overlaps(u):
532 return u.conj().transpose(0, 2, 1)[:, None] @ m0 @ u[neighbour]
534 def trial_gauge(u, d, x):
535 u_x = u @ _unitary_step(d, x)
536 m_x = overlaps(u_x)
537 c_x, s_x, *_ = _spread(m_x, b, w, n_k)
538 return float(np.sum(s_x)), u_x, m_x, c_x
540 m = overlaps(u)
541 centers, spreads, *_ = _spread(m, b, w, n_k)
542 history = [float(np.sum(spreads))]
543 eps = step / (4 * np.sum(w))
544 converged = max_iter == 0
545 d, g_norm = None, 0.
546 for it in range(max_iter):
547 # conjugate gradients (Fletcher-Reeves, restarted every 5 steps; a
548 # direction without descent fails the line search, which restarts)
549 grad = _gradient(m, b, w, centers)
550 g_new = float(np.sum(np.abs(grad) ** 2))
551 d = grad + g_new / g_norm * d if d is not None and it % 5 and g_norm > 0. else grad
552 g_norm = g_new
553 # parabola through steps 0, eps, 2 eps; keep the lowest spread found
554 tries = [trial_gauge(u, d, eps), trial_gauge(u, d, 2 * eps)]
555 o0, o1, o2 = history[-1], tries[0][0], tries[1][0]
556 curv = o2 - 2 * o1 + o0
557 if curv > 0.:
558 x = (3 * o0 - 4 * o1 + o2) / (2 * curv)
559 if 0. < x < 4.:
560 tries.append(trial_gauge(u, d, x * eps))
561 best = min(tries, key=lambda t: t[0])
562 if best[0] >= o0:
563 # no descent along d: restart from the gradient with a smaller step,
564 # until the step is negligible (a minimum, within roundoff)
565 eps /= 2
566 d = None
567 if eps < 1e-6 * step / (4 * np.sum(w)):
568 converged = True
569 break
570 continue
571 omega, u, m, centers = best
572 history.append(omega)
573 if o0 - omega < tol:
574 converged = True
575 break
576 centers, spreads, omega_i, omega_d, omega_od = _spread(m, b, w, n_k)
578 # Wannier functions: W_n(R, o) = (1/N) sum_k e^{ik.R} (v U)_{on}(k)
579 amp = np.fft.ifftn((v @ u).reshape(nk + (ks.norb, n_w)), axes=tuple(range(ks.dim)))
580 # move each one so that its centre sits in the middle cell
581 rec = np.array(ks.rec_vec, dtype='f8')
582 cart = centers @ ks.k_basis.T
583 frac = cart @ rec.T / (2 * np.pi)
584 shift = np.array(nk)[None, :] // 2 - np.floor(frac + 1e-8).astype(int) # roundoff at a cell corner
585 functions = np.zeros((n_w, n_k * ks.norb), 'c16')
586 for n in range(n_w):
587 moved = np.roll(amp[..., n], tuple(shift[n]), axis=tuple(range(ks.dim)))
588 # rows of finite_ham: cell n1 + N1 n2 + N1 N2 n3, then the orbital
589 functions[n] = np.transpose(moved, tuple(range(ks.dim))[::-1] + (ks.dim,)).ravel()
590 cart = cart + shift @ np.array(ks.lat.prim_vec, dtype='f8')
591 positions, _ = _finite_sites(ks, nk)
592 # out of the plane of a 2D lattice in 3D (or off a chain), the position
593 # operator is well defined: the centre is its expectation value
594 perp = np.eye(ks.space_dim) - ks.k_basis @ ks.k_basis.T
595 cart = cart + (np.abs(functions) ** 2 @ positions) @ perp
596 gauge = u_dis @ u # (n_k, len(bands), n_w)
597 ham_k = gauge.conj().transpose(0, 2, 1) @ (energies[:, :, None] * gauge)
598 return WannierFunctions(ks, nk, functions, positions, cart, spreads, omega_i, omega_d,
599 omega_od, gauge, np.array(history), converged, float(s.min()),
600 dis_history, dis_converged, ham_k, shift)