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