Coverage for /builds/kinetik161/ase/ase/dft/wannier.py: 58.84%
554 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-12-10 11:04 +0000
1""" Partly occupied Wannier functions
3 Find the set of partly occupied Wannier functions using the method from
4 Thygesen, Hansen and Jacobsen PRB v72 i12 p125119 2005.
5"""
6import functools
7import warnings
8from math import pi, sqrt
9from time import time
11import numpy as np
12from scipy.linalg import qr
14from ase.dft.bandgap import bandgap
15from ase.dft.kpoints import get_monkhorst_pack_size_and_offset
16from ase.io.jsonio import read_json, write_json
17from ase.parallel import paropen
18from ase.transport.tools import dagger, normalize
20dag = dagger
23def silent(*args, **kwargs):
24 """Dummy logging function."""
27def gram_schmidt(U):
28 """Orthonormalize columns of U according to the Gram-Schmidt procedure."""
29 for i, col in enumerate(U.T):
30 for col2 in U.T[:i]:
31 col -= col2 * (col2.conj() @ col)
32 col /= np.linalg.norm(col)
35def lowdin(U, S=None):
36 """Orthonormalize columns of U according to the symmetric Lowdin procedure.
37 The implementation uses SVD, like symm. Lowdin it returns the nearest
38 orthonormal matrix, but is more robust.
39 """
41 L, s, R = np.linalg.svd(U, full_matrices=False)
42 U[:] = L @ R
45def neighbor_k_search(k_c, G_c, kpt_kc, tol=1e-4):
46 # search for k1 (in kpt_kc) and k0 (in alldir), such that
47 # k1 - k - G + k0 = 0
48 alldir_dc = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1],
49 [1, 1, 0], [1, 0, 1], [0, 1, 1]], dtype=int)
50 for k0_c in alldir_dc:
51 for k1, k1_c in enumerate(kpt_kc):
52 if np.linalg.norm(k1_c - k_c - G_c + k0_c) < tol:
53 return k1, k0_c
55 raise ValueError(f'Wannier: Did not find matching kpoint for kpt={k_c}. '
56 'Probably non-uniform k-point grid')
59def calculate_weights(cell_cc, normalize=True):
60 """ Weights are used for non-cubic cells, see PRB **61**, 10040
61 If normalized they lose the physical dimension."""
62 alldirs_dc = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1],
63 [1, 1, 0], [1, 0, 1], [0, 1, 1]], dtype=int)
64 g = cell_cc @ cell_cc.T
65 # NOTE: Only first 3 of following 6 weights are presently used:
66 w = np.zeros(6)
67 w[0] = g[0, 0] - g[0, 1] - g[0, 2]
68 w[1] = g[1, 1] - g[0, 1] - g[1, 2]
69 w[2] = g[2, 2] - g[0, 2] - g[1, 2]
70 w[3] = g[0, 1]
71 w[4] = g[0, 2]
72 w[5] = g[1, 2]
73 # Make sure that first 3 Gdir vectors are included -
74 # these are used to calculate Wanniercenters.
75 Gdir_dc = alldirs_dc[:3]
76 weight_d = w[:3]
77 for d in range(3, 6):
78 if abs(w[d]) > 1e-5:
79 Gdir_dc = np.concatenate((Gdir_dc, alldirs_dc[d:d + 1]))
80 weight_d = np.concatenate((weight_d, w[d:d + 1]))
81 if normalize:
82 weight_d /= max(abs(weight_d))
83 return weight_d, Gdir_dc
86def steepest_descent(func, step=.005, tolerance=1e-6, log=silent, **kwargs):
87 fvalueold = 0.
88 fvalue = fvalueold + 10
89 count = 0
90 while abs((fvalue - fvalueold) / fvalue) > tolerance:
91 fvalueold = fvalue
92 dF = func.get_gradients()
93 func.step(dF * step, **kwargs)
94 fvalue = func.get_functional_value()
95 count += 1
96 log(f'SteepestDescent: iter={count}, value={fvalue}')
99def md_min(func, step=.25, tolerance=1e-6, max_iter=10000,
100 log=silent, **kwargs):
102 log('Localize with step =', step, 'and tolerance =', tolerance)
103 finit = func.get_functional_value()
105 t = -time()
106 fvalueold = 0.
107 fvalue = fvalueold + 10
108 count = 0
109 V = np.zeros(func.get_gradients().shape, dtype=complex)
111 while abs((fvalue - fvalueold) / fvalue) > tolerance:
112 fvalueold = fvalue
113 dF = func.get_gradients()
115 V *= (dF * V.conj()).real > 0
116 V += step * dF
117 func.step(V, **kwargs)
118 fvalue = func.get_functional_value()
120 if fvalue < fvalueold:
121 step *= 0.5
122 count += 1
123 log(f'MDmin: iter={count}, step={step}, value={fvalue}')
124 if count > max_iter:
125 t += time()
126 warnings.warn('Max iterations reached: '
127 'iters=%s, step=%s, seconds=%0.2f, value=%0.4f'
128 % (count, step, t, fvalue.real))
129 break
131 t += time()
132 log('%d iterations in %0.2f seconds (%0.2f ms/iter), endstep = %s' %
133 (count, t, t * 1000. / count, step))
134 log(f'Initial value={finit}, Final value={fvalue}')
137def rotation_from_projection(proj_nw, fixed, ortho=True):
138 """Determine rotation and coefficient matrices from projections
140 proj_nw = <psi_n|p_w>
141 psi_n: eigenstates
142 p_w: localized function
144 Nb (n) = Number of bands
145 Nw (w) = Number of wannier functions
146 M (f) = Number of fixed states
147 L (l) = Number of extra degrees of freedom
148 U (u) = Number of non-fixed states
149 """
151 Nb, Nw = proj_nw.shape
152 M = fixed
153 L = Nw - M
154 U = Nb - M
156 U_ww = np.empty((Nw, Nw), dtype=proj_nw.dtype)
158 # Set the section of the rotation matrix about the 'fixed' states
159 U_ww[:M] = proj_nw[:M]
161 if L > 0:
162 # If there are extra degrees of freedom we have to select L of them
163 C_ul = np.empty((U, L), dtype=proj_nw.dtype)
165 # Get the projections on the 'non fixed' states
166 proj_uw = proj_nw[M:]
168 # Obtain eigenvalues and eigevectors matrix
169 eig_w, C_ww = np.linalg.eigh(dag(proj_uw) @ proj_uw)
171 # Sort columns of eigenvectors matrix according to the eigenvalues
172 # magnitude, select only the L largest ones. Then use them to obtain
173 # the parameter C matrix.
174 C_ul[:] = proj_uw @ C_ww[:, np.argsort(- eig_w.real)[:L]]
176 # Compute the section of the rotation matrix about 'non fixed' states
177 U_ww[M:] = dag(C_ul) @ proj_uw
178 normalize(C_ul)
179 else:
180 # If there are no extra degrees of freedom we do not need any parameter
181 # matrix C
182 C_ul = np.empty((U, 0), dtype=proj_nw.dtype)
184 if ortho:
185 # Orthogonalize with Lowdin to take the closest orthogonal set
186 lowdin(U_ww)
187 else:
188 normalize(U_ww)
190 return U_ww, C_ul
193def search_for_gamma_point(kpts):
194 """Returns index of Gamma point in a list of k-points."""
195 gamma_idx = np.argmin([np.linalg.norm(kpt) for kpt in kpts])
196 if not np.linalg.norm(kpts[gamma_idx]) < 1e-14:
197 gamma_idx = None
198 return gamma_idx
201def scdm(pseudo_nkG, kpts, fixed_k, Nw):
202 """Compute localized orbitals with SCDM method
204 This method was published by Anil Damle and Lin Lin in Multiscale
205 Modeling & Simulation 16, 1392–1410 (2018).
206 For now only the isolated bands algorithm is implemented, because it is
207 intended as a drop-in replacement for other initial guess methods for
208 the ASE Wannier class.
210 pseudo_nkG = pseudo wave-functions on a real grid
211 Ng (G) = number of real grid points
212 kpts = List of k-points in the BZ
213 Nk (k) = Number of k-points
214 Nb (n) = Number of bands
215 Nw (w) = Number of wannier functions
216 fixed_k = Number of fixed states for each k-point
217 L (l) = Number of extra degrees of freedom
218 U (u) = Number of non-fixed states
219 """
221 gamma_idx = search_for_gamma_point(kpts)
222 Nk = len(kpts)
223 U_kww = []
224 C_kul = []
226 # compute factorization only at Gamma point
227 _, _, P = qr(pseudo_nkG[:, gamma_idx, :], mode='full',
228 pivoting=True, check_finite=True)
230 for k in range(Nk):
231 A_nw = pseudo_nkG[:, k, P[:Nw]]
232 U_ww, C_ul = rotation_from_projection(proj_nw=A_nw,
233 fixed=fixed_k[k],
234 ortho=True)
235 U_kww.append(U_ww)
236 C_kul.append(C_ul)
238 U_kww = np.asarray(U_kww)
240 return C_kul, U_kww
243def arbitrary_s_orbitals(atoms, Ns, rng=np.random):
244 """
245 Generate a list of Ns randomly placed s-orbitals close to at least
246 one atom (< 1.5Å).
247 The format of the list is the one required by GPAW in initial_wannier().
248 """
249 # Create dummy copy of the Atoms object and dummy H atom
250 tmp_atoms = atoms.copy()
251 tmp_atoms.append('H')
252 s_pos = tmp_atoms.get_scaled_positions()
254 orbs = []
255 for i in range(0, Ns):
256 fine = False
257 while not fine:
258 # Random position
259 x, y, z = rng.rand(3)
260 s_pos[-1] = [x, y, z]
261 tmp_atoms.set_scaled_positions(s_pos)
263 # Use dummy H atom to measure distance from any other atom
264 dists = tmp_atoms.get_distances(
265 a=-1,
266 indices=range(len(atoms)))
268 # Check if it is close to at least one atom
269 if (dists < 1.5).any():
270 fine = True
272 orbs.append([[x, y, z], 0, 1])
273 return orbs
276def init_orbitals(atoms, ntot, rng=np.random):
277 """
278 Place d-orbitals for every atom that has some in the valence states
279 and then random s-orbitals close to at least one atom (< 1.5Å).
280 The orbitals list format is compatible with GPAW.get_initial_wannier().
281 'atoms': ASE Atoms object
282 'ntot': total number of needed orbitals
283 'rng': generator random numbers
284 """
286 # List all the elements that should have occupied d-orbitals
287 # in the valence states (according to GPAW setups)
288 d_metals = set(list(range(21, 31)) + list(range(39, 52)) +
289 list(range(57, 84)) + list(range(89, 113)))
290 orbs = []
292 # Start with zero orbitals
293 No = 0
295 # Add d orbitals to each d-metal
296 for i, z in enumerate(atoms.get_atomic_numbers()):
297 if z in d_metals:
298 No_new = No + 5
299 if No_new <= ntot:
300 orbs.append([i, 2, 1])
301 No = No_new
303 if No < ntot:
304 # Add random s-like orbitals if there are not enough yet
305 Ns = ntot - No
306 orbs += arbitrary_s_orbitals(atoms, Ns, rng)
308 assert sum([orb[1] * 2 + 1 for orb in orbs]) == ntot
309 return orbs
312def square_modulus_of_Z_diagonal(Z_dww):
313 """
314 Square modulus of the Z matrix diagonal, the diagonal is taken
315 for the indexes running on the WFs.
316 """
317 return np.abs(Z_dww.diagonal(0, 1, 2))**2
320def get_kklst(kpt_kc, Gdir_dc):
321 # Set the list of neighboring k-points k1, and the "wrapping" k0,
322 # such that k1 - k - G + k0 = 0
323 #
324 # Example: kpoints = (-0.375,-0.125,0.125,0.375), dir=0
325 # G = [0.25,0,0]
326 # k=0.375, k1= -0.375 : -0.375-0.375-0.25 => k0=[1,0,0]
327 #
328 # For a gamma point calculation k1 = k = 0, k0 = [1,0,0] for dir=0
329 Nk = len(kpt_kc)
330 Ndir = len(Gdir_dc)
332 if Nk == 1:
333 kklst_dk = np.zeros((Ndir, 1), int)
334 k0_dkc = Gdir_dc.reshape(-1, 1, 3)
335 else:
336 kklst_dk = np.empty((Ndir, Nk), int)
337 k0_dkc = np.empty((Ndir, Nk, 3), int)
339 # Distance between kpoints
340 kdist_c = np.empty(3)
341 for c in range(3):
342 # make a sorted list of the kpoint values in this direction
343 slist = np.argsort(kpt_kc[:, c], kind='mergesort')
344 skpoints_kc = np.take(kpt_kc, slist, axis=0)
345 kdist_c[c] = max([skpoints_kc[n + 1, c] - skpoints_kc[n, c]
346 for n in range(Nk - 1)])
348 for d, Gdir_c in enumerate(Gdir_dc):
349 for k, k_c in enumerate(kpt_kc):
350 # setup dist vector to next kpoint
351 G_c = np.where(Gdir_c > 0, kdist_c, 0)
352 if max(G_c) < 1e-4:
353 kklst_dk[d, k] = k
354 k0_dkc[d, k] = Gdir_c
355 else:
356 kklst_dk[d, k], k0_dkc[d, k] = \
357 neighbor_k_search(k_c, G_c, kpt_kc)
358 return kklst_dk, k0_dkc
361def get_invkklst(kklst_dk):
362 Ndir, Nk = kklst_dk.shape
363 invkklst_dk = np.empty(kklst_dk.shape, int)
364 for d in range(Ndir):
365 for k1 in range(Nk):
366 invkklst_dk[d, k1] = kklst_dk[d].tolist().index(k1)
367 return invkklst_dk
370def choose_states(calcdata, fixedenergy, fixedstates, Nk, nwannier, log, spin):
372 if fixedenergy is None and fixedstates is not None:
373 if isinstance(fixedstates, int):
374 fixedstates = [fixedstates] * Nk
375 fixedstates_k = np.array(fixedstates, int)
376 elif fixedenergy is not None and fixedstates is None:
377 # Setting number of fixed states and EDF from given energy cutoff.
378 # All states below this energy cutoff are fixed.
379 # The reference energy is Ef for metals and CBM for insulators.
380 if calcdata.gap < 0.01 or fixedenergy < 0.01:
381 cutoff = fixedenergy + calcdata.fermi_level
382 else:
383 cutoff = fixedenergy + calcdata.lumo
385 # Find the states below the energy cutoff at each k-point
386 tmp_fixedstates_k = []
387 for k in range(Nk):
388 eps_n = calcdata.eps_skn[spin, k]
389 kindex = eps_n.searchsorted(cutoff)
390 tmp_fixedstates_k.append(kindex)
391 fixedstates_k = np.array(tmp_fixedstates_k, int)
392 elif fixedenergy is not None and fixedstates is not None:
393 raise RuntimeError(
394 'You can not set both fixedenergy and fixedstates')
396 if nwannier == 'auto':
397 if fixedenergy is None and fixedstates is None:
398 # Assume the fixedexergy parameter equal to 0 and
399 # find the states below the Fermi level at each k-point.
400 log("nwannier=auto but no 'fixedenergy' or 'fixedstates'",
401 "parameter was provided, using Fermi level as",
402 "energy cutoff.")
403 tmp_fixedstates_k = []
404 for k in range(Nk):
405 eps_n = calcdata.eps_skn[spin, k]
406 kindex = eps_n.searchsorted(calcdata.fermi_level)
407 tmp_fixedstates_k.append(kindex)
408 fixedstates_k = np.array(tmp_fixedstates_k, int)
409 nwannier = np.max(fixedstates_k)
411 # Without user choice just set nwannier fixed states without EDF
412 if fixedstates is None and fixedenergy is None:
413 fixedstates_k = np.array([nwannier] * Nk, int)
415 return fixedstates_k, nwannier
418def get_eigenvalues(calc):
419 nspins = calc.get_number_of_spins()
420 nkpts = len(calc.get_ibz_k_points())
421 nbands = calc.get_number_of_bands()
422 eps_skn = np.empty((nspins, nkpts, nbands))
424 for ispin in range(nspins):
425 for ikpt in range(nkpts):
426 eps_skn[ispin, ikpt] = calc.get_eigenvalues(kpt=ikpt, spin=ispin)
427 return eps_skn
430class CalcData:
431 def __init__(self, kpt_kc, atoms, fermi_level, lumo, eps_skn,
432 gap):
433 self.kpt_kc = kpt_kc
434 self.atoms = atoms
435 self.fermi_level = fermi_level
436 self.lumo = lumo
437 self.eps_skn = eps_skn
438 self.gap = gap
440 @property
441 def nbands(self):
442 return self.eps_skn.shape[2]
445def get_calcdata(calc):
446 kpt_kc = calc.get_bz_k_points()
447 # Make sure there is no symmetry reduction
448 assert len(calc.get_ibz_k_points()) == len(kpt_kc)
449 lumo = calc.get_homo_lumo()[1]
450 gap = bandgap(calc=calc, output=None)[0]
451 return CalcData(
452 kpt_kc=kpt_kc,
453 atoms=calc.get_atoms(),
454 fermi_level=calc.get_fermi_level(),
455 lumo=lumo,
456 eps_skn=get_eigenvalues(calc),
457 gap=gap)
460class Wannier:
461 """Partly occupied Wannier functions
463 Find the set of partly occupied Wannier functions according to
464 Thygesen, Hansen and Jacobsen PRB v72 i12 p125119 2005.
465 """
467 def __init__(self, nwannier, calc,
468 file=None,
469 nbands=None,
470 fixedenergy=None,
471 fixedstates=None,
472 spin=0,
473 initialwannier='orbitals',
474 functional='std',
475 rng=np.random,
476 log=silent):
477 """
478 Required arguments:
480 ``nwannier``: The number of Wannier functions you wish to construct.
481 This must be at least half the number of electrons in the system
482 and at most equal to the number of bands in the calculation.
483 It can also be set to 'auto' in order to automatically choose the
484 minimum number of needed Wannier function based on the
485 ``fixedenergy`` / ``fixedstates`` parameter.
487 ``calc``: A converged DFT calculator class.
488 If ``file`` arg. is not provided, the calculator *must* provide the
489 method ``get_wannier_localization_matrix``, and contain the
490 wavefunctions (save files with only the density is not enough).
491 If the localization matrix is read from file, this is not needed,
492 unless ``get_function`` or ``write_cube`` is called.
494 Optional arguments:
496 ``nbands``: Bands to include in localization.
497 The number of bands considered by Wannier can be smaller than the
498 number of bands in the calculator. This is useful if the highest
499 bands of the DFT calculation are not well converged.
501 ``spin``: The spin channel to be considered.
502 The Wannier code treats each spin channel independently.
504 ``fixedenergy`` / ``fixedstates``: Fixed part of Hilbert space.
505 Determine the fixed part of Hilbert space by either a maximal
506 energy *or* a number of bands (possibly a list for multiple
507 k-points).
508 Default is None meaning that the number of fixed states is equated
509 to ``nwannier``.
510 The maximal energy is relative to the CBM if there is a finite
511 bandgap or to the Fermi level if there is none.
513 ``file``: Read localization and rotation matrices from this file.
515 ``initialwannier``: Initial guess for Wannier rotation matrix.
516 Can be 'bloch' to start from the Bloch states, 'random' to be
517 randomized, 'orbitals' to start from atom-centered d-orbitals and
518 randomly placed gaussian centers (see init_orbitals()),
519 'scdm' to start from localized state selected with SCDM
520 or a list passed to calc.get_initial_wannier.
522 ``functional``: The functional used to measure the localization.
523 Can be 'std' for the standard quadratic functional from the PRB
524 paper, 'var' to add a variance minimizing term.
526 ``rng``: Random number generator for ``initialwannier``.
528 ``log``: Function which logs, such as print().
529 """
530 # Bloch phase sign convention.
531 # May require special cases depending on which code is used.
532 sign = -1
534 self.log = log
535 self.calc = calc
537 self.spin = spin
538 self.functional = functional
539 self.initialwannier = initialwannier
540 self.log('Using functional:', functional)
542 self.calcdata = get_calcdata(calc)
544 self.kptgrid = get_monkhorst_pack_size_and_offset(self.kpt_kc)[0]
545 self.calcdata.kpt_kc *= sign
547 self.largeunitcell_cc = (self.unitcell_cc.T * self.kptgrid).T
548 self.weight_d, self.Gdir_dc = calculate_weights(self.largeunitcell_cc)
549 assert len(self.weight_d) == len(self.Gdir_dc)
551 if nbands is None:
552 # XXX Can work with other number of bands than calculator.
553 # Is this case properly tested, lest we confuse them?
554 nbands = self.calcdata.nbands
555 self.nbands = nbands
557 self.fixedstates_k, self.nwannier = choose_states(
558 self.calcdata, fixedenergy, fixedstates, self.Nk, nwannier,
559 log, spin)
561 # Compute the number of extra degrees of freedom (EDF)
562 self.edf_k = self.nwannier - self.fixedstates_k
564 self.log(f'Wannier: Fixed states : {self.fixedstates_k}')
565 self.log(f'Wannier: Extra degrees of freedom: {self.edf_k}')
567 self.kklst_dk, k0_dkc = get_kklst(self.kpt_kc, self.Gdir_dc)
569 # Set the inverse list of neighboring k-points
570 self.invkklst_dk = get_invkklst(self.kklst_dk)
572 Nw = self.nwannier
573 Nb = self.nbands
574 self.Z_dkww = np.empty((self.Ndir, self.Nk, Nw, Nw), complex)
575 self.V_knw = np.zeros((self.Nk, Nb, Nw), complex)
577 if file is None:
578 self.Z_dknn = self.new_Z(calc, k0_dkc)
579 self.initialize(file=file, initialwannier=initialwannier, rng=rng)
581 @property
582 def atoms(self):
583 return self.calcdata.atoms
585 @property
586 def kpt_kc(self):
587 return self.calcdata.kpt_kc
589 @property
590 def Ndir(self):
591 return len(self.weight_d) # Number of directions
593 @property
594 def Nk(self):
595 return len(self.kpt_kc)
597 def new_Z(self, calc, k0_dkc):
598 Nb = self.nbands
599 Z_dknn = np.empty((self.Ndir, self.Nk, Nb, Nb), complex)
600 for d, dirG in enumerate(self.Gdir_dc):
601 for k in range(self.Nk):
602 k1 = self.kklst_dk[d, k]
603 k0_c = k0_dkc[d, k]
604 Z_dknn[d, k] = calc.get_wannier_localization_matrix(
605 nbands=Nb, dirG=dirG, kpoint=k, nextkpoint=k1,
606 G_I=k0_c, spin=self.spin)
607 return Z_dknn
609 @property
610 def unitcell_cc(self):
611 return self.atoms.cell
613 @property
614 def U_kww(self):
615 return self.wannier_state.U_kww
617 @property
618 def C_kul(self):
619 return self.wannier_state.C_kul
621 def initialize(self, file=None, initialwannier='random', rng=np.random):
622 """Re-initialize current rotation matrix.
624 Keywords are identical to those of the constructor.
625 """
626 from ase.dft.wannierstate import WannierSpec, WannierState
628 spec = WannierSpec(self.Nk, self.nwannier, self.nbands,
629 self.fixedstates_k)
631 if file is not None:
632 with paropen(file, 'r') as fd:
633 Z_dknn, U_kww, C_kul = read_json(fd, always_array=False)
634 self.Z_dknn = Z_dknn
635 wannier_state = WannierState(C_kul, U_kww)
636 elif initialwannier == 'bloch':
637 # Set U and C to pick the lowest Bloch states
638 wannier_state = spec.bloch(self.edf_k)
639 elif initialwannier == 'random':
640 wannier_state = spec.random(rng, self.edf_k)
641 elif initialwannier == 'orbitals':
642 orbitals = init_orbitals(self.atoms, self.nwannier, rng)
643 wannier_state = spec.initial_orbitals(
644 self.calc, orbitals, self.kptgrid, self.edf_k, self.spin)
645 elif initialwannier == 'scdm':
646 wannier_state = spec.scdm(self.calc, self.kpt_kc, self.spin)
647 else:
648 wannier_state = spec.initial_wannier(
649 self.calc, initialwannier, self.kptgrid,
650 self.edf_k, self.spin)
652 self.wannier_state = wannier_state
653 self.update()
655 def save(self, file):
656 """Save information on localization and rotation matrices to file."""
657 with paropen(file, 'w') as fd:
658 write_json(fd, (self.Z_dknn, self.U_kww, self.C_kul))
660 def update(self):
661 # Update large rotation matrix V (from rotation U and coeff C)
662 for k, M in enumerate(self.fixedstates_k):
663 self.V_knw[k, :M] = self.U_kww[k, :M]
664 if M < self.nwannier:
665 self.V_knw[k, M:] = self.C_kul[k] @ self.U_kww[k, M:]
666 # else: self.V_knw[k, M:] = 0.0
668 # Calculate the Zk matrix from the large rotation matrix:
669 # Zk = V^d[k] Zbloch V[k1]
670 for d in range(self.Ndir):
671 for k in range(self.Nk):
672 k1 = self.kklst_dk[d, k]
673 self.Z_dkww[d, k] = dag(self.V_knw[k]) \
674 @ (self.Z_dknn[d, k] @ self.V_knw[k1])
676 # Update the new Z matrix
677 self.Z_dww = self.Z_dkww.sum(axis=1) / self.Nk
679 def get_optimal_nwannier(self, nwrange=5, random_reps=5, tolerance=1e-6):
680 """
681 The optimal value for 'nwannier', maybe.
683 The optimal value is the one that gives the lowest average value for
684 the spread of the most delocalized Wannier function in the set.
686 ``nwrange``: number of different values to try for 'nwannier', the
687 values will span a symmetric range around ``nwannier`` if possible.
689 ``random_reps``: number of repetitions with random seed, the value is
690 then an average over these repetitions.
692 ``tolerance``: tolerance for the gradient descent algorithm, can be
693 useful to increase the speed, with a cost in accuracy.
694 """
696 # Define the range of values to try based on the maximum number of fixed
697 # states (that is the minimum number of WFs we need) and the number of
698 # available bands we have.
699 max_number_fixedstates = np.max(self.fixedstates_k)
701 min_range_value = max(self.nwannier - int(np.floor(nwrange / 2)),
702 max_number_fixedstates)
703 max_range_value = min(min_range_value + nwrange, self.nbands + 1)
704 Nws = np.arange(min_range_value, max_range_value)
706 # If there is no randomness, there is no need to repeat
707 random_initials = ['random', 'orbitals']
708 if self.initialwannier not in random_initials:
709 random_reps = 1
711 t = -time()
712 avg_max_spreads = np.zeros(len(Nws))
713 for j, Nw in enumerate(Nws):
714 self.log('Trying with Nw =', Nw)
716 # Define once with the fastest 'initialwannier',
717 # then initialize with random seeds in the for loop
718 wan = Wannier(nwannier=int(Nw),
719 calc=self.calc,
720 nbands=self.nbands,
721 spin=self.spin,
722 functional=self.functional,
723 initialwannier='bloch',
724 log=self.log,
725 rng=self.rng)
726 wan.fixedstates_k = self.fixedstates_k
727 wan.edf_k = wan.nwannier - wan.fixedstates_k
729 max_spreads = np.zeros(random_reps)
730 for i in range(random_reps):
731 wan.initialize(initialwannier=self.initialwannier)
732 wan.localize(tolerance=tolerance)
733 max_spreads[i] = np.max(wan.get_spreads())
735 avg_max_spreads[j] = max_spreads.mean()
737 self.log('Average spreads: ', avg_max_spreads)
738 t += time()
739 self.log(f'Execution time: {t:.1f}s')
741 return Nws[np.argmin(avg_max_spreads)]
743 def get_centers(self, scaled=False):
744 """Calculate the Wannier centers
746 ::
748 pos = L / 2pi * phase(diag(Z))
749 """
750 coord_wc = np.angle(self.Z_dww[:3].diagonal(0, 1, 2)).T / (2 * pi) % 1
751 if not scaled:
752 coord_wc = coord_wc @ self.largeunitcell_cc
753 return coord_wc
755 def get_radii(self):
756 r"""Calculate the spread of the Wannier functions.
758 ::
760 -- / L \ 2 2
761 radius**2 = - > | --- | ln |Z|
762 --d \ 2pi /
764 Note that this function can fail with some Bravais lattices,
765 see `get_spreads()` for a more robust alternative.
766 """
767 r2 = - (self.largeunitcell_cc.diagonal()**2 / (2 * pi)**2) \
768 @ np.log(abs(self.Z_dww[:3].diagonal(0, 1, 2))**2)
769 return np.sqrt(r2)
771 def get_spreads(self):
772 r"""Calculate the spread of the Wannier functions in Ų.
773 The definition is based on eq. 13 in PRB61-15 by Berghold and Mundy.
775 ::
777 / 1 \ 2 -- 2
778 spread = - |----| > W_d ln |Z_dw|
779 \2 pi/ --d
782 """
783 # compute weights without normalization, to keep physical dimension
784 weight_d, _ = calculate_weights(self.largeunitcell_cc, normalize=False)
785 Z2_dw = square_modulus_of_Z_diagonal(self.Z_dww)
786 spread_w = - (np.log(Z2_dw).T @ weight_d).real / (2 * np.pi)**2
787 return spread_w
789 def get_spectral_weight(self, w):
790 return abs(self.V_knw[:, :, w])**2 / self.Nk
792 def get_pdos(self, w, energies, width):
793 """Projected density of states (PDOS).
795 Returns the (PDOS) for Wannier function ``w``. The calculation
796 is performed over the energy grid specified in energies. The
797 PDOS is produced as a sum of Gaussians centered at the points
798 of the energy grid and with the specified width.
799 """
800 spec_kn = self.get_spectral_weight(w)
801 dos = np.zeros(len(energies))
802 for k, spec_n in enumerate(spec_kn):
803 eig_n = self.calcdata.eps_skn[self.spin, k]
804 for weight, eig in zip(spec_n, eig_n):
805 # Add gaussian centered at the eigenvalue
806 x = ((energies - eig) / width)**2
807 dos += weight * np.exp(-x.clip(0., 40.)) / (sqrt(pi) * width)
808 return dos
810 def translate(self, w, R):
811 """Translate the w'th Wannier function
813 The distance vector R = [n1, n2, n3], is in units of the basis
814 vectors of the small cell.
815 """
816 for kpt_c, U_ww in zip(self.kpt_kc, self.U_kww):
817 U_ww[:, w] *= np.exp(2.j * pi * (np.array(R) @ kpt_c))
818 self.update()
820 def translate_to_cell(self, w, cell):
821 """Translate the w'th Wannier function to specified cell"""
822 scaled_c = np.angle(self.Z_dww[:3, w, w]) * self.kptgrid / (2 * pi)
823 trans = np.array(cell) - np.floor(scaled_c)
824 self.translate(w, trans)
826 def translate_all_to_cell(self, cell=(0, 0, 0)):
827 r"""Translate all Wannier functions to specified cell.
829 Move all Wannier orbitals to a specific unit cell. There
830 exists an arbitrariness in the positions of the Wannier
831 orbitals relative to the unit cell. This method can move all
832 orbitals to the unit cell specified by ``cell``. For a
833 `\Gamma`-point calculation, this has no effect. For a
834 **k**-point calculation the periodicity of the orbitals are
835 given by the large unit cell defined by repeating the original
836 unitcell by the number of **k**-points in each direction. In
837 this case it is useful to move the orbitals away from the
838 boundaries of the large cell before plotting them. For a bulk
839 calculation with, say 10x10x10 **k** points, one could move
840 the orbitals to the cell [2,2,2]. In this way the pbc
841 boundary conditions will not be noticed.
842 """
843 scaled_wc = (np.angle(self.Z_dww[:3].diagonal(0, 1, 2)).T *
844 self.kptgrid / (2 * pi))
845 trans_wc = np.array(cell)[None] - np.floor(scaled_wc)
846 for kpt_c, U_ww in zip(self.kpt_kc, self.U_kww):
847 U_ww *= np.exp(2.j * pi * (trans_wc @ kpt_c))
848 self.update()
850 def distances(self, R):
851 """Relative distances between centers.
853 Returns a matrix with the distances between different Wannier centers.
854 R = [n1, n2, n3] is in units of the basis vectors of the small cell
855 and allows one to measure the distance with centers moved to a
856 different small cell.
857 The dimension of the matrix is [Nw, Nw].
858 """
859 Nw = self.nwannier
860 cen = self.get_centers()
861 r1 = cen.repeat(Nw, axis=0).reshape(Nw, Nw, 3)
862 r2 = cen.copy()
863 for i in range(3):
864 r2 += self.unitcell_cc[i] * R[i]
866 r2 = np.swapaxes(r2.repeat(Nw, axis=0).reshape(Nw, Nw, 3), 0, 1)
867 return np.sqrt(np.sum((r1 - r2)**2, axis=-1))
869 @functools.lru_cache(maxsize=10000)
870 def _get_hopping(self, n1, n2, n3):
871 """Returns the matrix H(R)_nm=<0,n|H|R,m>.
873 ::
875 1 _ -ik.R
876 H(R) = <0,n|H|R,m> = --- >_ e H(k)
877 Nk k
879 where R = (n1, n2, n3) is the cell-distance (in units of the basis
880 vectors of the small cell) and n,m are indices of the Wannier functions.
881 This function caches up to 'maxsize' results.
882 """
883 R = np.array([n1, n2, n3], float)
884 H_ww = np.zeros([self.nwannier, self.nwannier], complex)
885 for k, kpt_c in enumerate(self.kpt_kc):
886 phase = np.exp(-2.j * pi * (np.array(R) @ kpt_c))
887 H_ww += self.get_hamiltonian(k) * phase
888 return H_ww / self.Nk
890 def get_hopping(self, R):
891 """Returns the matrix H(R)_nm=<0,n|H|R,m>.
893 ::
895 1 _ -ik.R
896 H(R) = <0,n|H|R,m> = --- >_ e H(k)
897 Nk k
899 where R is the cell-distance (in units of the basis vectors of
900 the small cell) and n,m are indices of the Wannier functions.
901 """
902 return self._get_hopping(R[0], R[1], R[2])
904 def get_hamiltonian(self, k):
905 """Get Hamiltonian at existing k-vector of index k
907 ::
909 dag
910 H(k) = V diag(eps ) V
911 k k k
912 """
913 eps_n = self.calcdata.eps_skn[self.spin, k, :self.nbands]
914 V_nw = self.V_knw[k]
915 return (dag(V_nw) * eps_n) @ V_nw
917 def get_hamiltonian_kpoint(self, kpt_c):
918 """Get Hamiltonian at some new arbitrary k-vector
920 ::
922 _ ik.R
923 H(k) = >_ e H(R)
924 R
926 Warning: This method moves all Wannier functions to cell (0, 0, 0)
927 """
928 self.log('Translating all Wannier functions to cell (0, 0, 0)')
929 self.translate_all_to_cell()
930 max = (self.kptgrid - 1) // 2
931 N1, N2, N3 = max
932 Hk = np.zeros([self.nwannier, self.nwannier], complex)
933 for n1 in range(-N1, N1 + 1):
934 for n2 in range(-N2, N2 + 1):
935 for n3 in range(-N3, N3 + 1):
936 R = np.array([n1, n2, n3], float)
937 hop_ww = self.get_hopping(R)
938 phase = np.exp(+2.j * pi * (R @ kpt_c))
939 Hk += hop_ww * phase
940 return Hk
942 def get_function(self, index, repeat=None):
943 r"""Get Wannier function on grid.
945 Returns an array with the funcion values of the indicated Wannier
946 function on a grid with the size of the *repeated* unit cell.
948 For a calculation using **k**-points the relevant unit cell for
949 eg. visualization of the Wannier orbitals is not the original unit
950 cell, but rather a larger unit cell defined by repeating the
951 original unit cell by the number of **k**-points in each direction.
952 Note that for a `\Gamma`-point calculation the large unit cell
953 coinsides with the original unit cell.
954 The large unitcell also defines the periodicity of the Wannier
955 orbitals.
957 ``index`` can be either a single WF or a coordinate vector in terms
958 of the WFs.
959 """
961 # Default size of plotting cell is the one corresponding to k-points.
962 if repeat is None:
963 repeat = self.kptgrid
964 N1, N2, N3 = repeat
966 dim = self.calc.get_number_of_grid_points()
967 largedim = dim * [N1, N2, N3]
969 wanniergrid = np.zeros(largedim, dtype=complex)
970 for k, kpt_c in enumerate(self.kpt_kc):
971 # The coordinate vector of wannier functions
972 if isinstance(index, int):
973 vec_n = self.V_knw[k, :, index]
974 else:
975 vec_n = self.V_knw[k] @ index
977 wan_G = np.zeros(dim, complex)
978 for n, coeff in enumerate(vec_n):
979 wan_G += coeff * self.calc.get_pseudo_wave_function(
980 n, k, self.spin, pad=True)
982 # Distribute the small wavefunction over large cell:
983 for n1 in range(N1):
984 for n2 in range(N2):
985 for n3 in range(N3): # sign?
986 e = np.exp(-2.j * pi * np.array([n1, n2, n3]) @ kpt_c)
987 wanniergrid[n1 * dim[0]:(n1 + 1) * dim[0],
988 n2 * dim[1]:(n2 + 1) * dim[1],
989 n3 * dim[2]:(n3 + 1) * dim[2]] += e * wan_G
991 # Normalization
992 wanniergrid /= np.sqrt(self.Nk)
993 return wanniergrid
995 def write_cube(self, index, fname, repeat=None, angle=False):
996 """
997 Dump specified Wannier function to a cube file.
999 Arguments:
1001 ``index``: Integer, index of the Wannier function to save.
1003 ``repeat``: Array of integer, repeat supercell and Wannier function.
1005 ``fname``: Name of the cube file.
1007 ``angle``: If False, save the absolute value. If True, save
1008 the complex phase of the Wannier function.
1009 """
1010 from ase.io import write
1012 # Default size of plotting cell is the one corresponding to k-points.
1013 if repeat is None:
1014 repeat = self.kptgrid
1016 # Remove constraints, some are not compatible with repeat()
1017 atoms = self.atoms.copy()
1018 atoms.set_constraint()
1019 atoms = atoms * repeat
1020 func = self.get_function(index, repeat)
1022 # Compute absolute value or complex angle
1023 if angle:
1024 data = np.angle(func)
1025 else:
1026 if self.Nk == 1:
1027 func *= np.exp(-1.j * np.angle(func.max()))
1028 func = abs(func)
1029 data = func
1031 write(fname, atoms, data=data, format='cube')
1033 def localize(self, step=0.25, tolerance=1e-08,
1034 updaterot=True, updatecoeff=True):
1035 """Optimize rotation to give maximal localization"""
1036 md_min(self, step=step, tolerance=tolerance, log=self.log,
1037 updaterot=updaterot, updatecoeff=updatecoeff)
1039 def get_functional_value(self):
1040 """Calculate the value of the spread functional.
1042 ::
1044 Tr[|ZI|^2]=sum(I)sum(n) w_i|Z_(i)_nn|^2,
1046 where w_i are weights.
1048 If the functional is set to 'var' it subtracts a variance term
1050 ::
1052 Nw * var(sum(n) w_i|Z_(i)_nn|^2),
1054 where Nw is the number of WFs ``nwannier``.
1055 """
1056 a_w = self._spread_contributions()
1057 if self.functional == 'std':
1058 fun = np.sum(a_w)
1059 elif self.functional == 'var':
1060 fun = np.sum(a_w) - self.nwannier * np.var(a_w)
1061 self.log(f'std: {np.sum(a_w):.4f}',
1062 f'\tvar: {self.nwannier * np.var(a_w):.4f}')
1063 return fun
1065 def get_gradients(self):
1066 # Determine gradient of the spread functional.
1067 #
1068 # The gradient for a rotation A_kij is::
1069 #
1070 # dU = dRho/dA_{k,i,j} = sum(I) sum(k')
1071 # + Z_jj Z_kk',ij^* - Z_ii Z_k'k,ij^*
1072 # - Z_ii^* Z_kk',ji + Z_jj^* Z_k'k,ji
1073 #
1074 # The gradient for a change of coefficients is::
1075 #
1076 # dRho/da^*_{k,i,j} = sum(I) [[(Z_0)_{k} V_{k'} diag(Z^*) +
1077 # (Z_0_{k''})^d V_{k''} diag(Z)] *
1078 # U_k^d]_{N+i,N+j}
1079 #
1080 # where diag(Z) is a square,diagonal matrix with Z_nn in the diagonal,
1081 # k' = k + dk and k = k'' + dk.
1082 #
1083 # The extra degrees of freedom chould be kept orthonormal to the fixed
1084 # space, thus we introduce lagrange multipliers, and minimize instead::
1085 #
1086 # Rho_L = Rho - sum_{k,n,m} lambda_{k,nm} <c_{kn}|c_{km}>
1087 #
1088 # for this reason the coefficient gradients should be multiplied
1089 # by (1 - c c^dag).
1091 Nb = self.nbands
1092 Nw = self.nwannier
1094 if self.functional == 'var':
1095 O_w = self._spread_contributions()
1096 O_sum = np.sum(O_w)
1098 dU = []
1099 dC = []
1100 for k in range(self.Nk):
1101 M = self.fixedstates_k[k]
1102 L = self.edf_k[k]
1103 U_ww = self.U_kww[k]
1104 C_ul = self.C_kul[k]
1105 Utemp_ww = np.zeros((Nw, Nw), complex)
1106 Ctemp_nw = np.zeros((Nb, Nw), complex)
1108 for d, weight in enumerate(self.weight_d):
1109 if abs(weight) < 1.0e-6:
1110 continue
1112 Z_knn = self.Z_dknn[d]
1113 diagZ_w = self.Z_dww[d].diagonal()
1114 Zii_ww = np.repeat(diagZ_w, Nw).reshape(Nw, Nw)
1115 if self.functional == 'var':
1116 diagOZ_w = O_w * diagZ_w
1117 OZii_ww = np.repeat(diagOZ_w, Nw).reshape(Nw, Nw)
1119 k1 = self.kklst_dk[d, k]
1120 k2 = self.invkklst_dk[d, k]
1121 V_knw = self.V_knw
1122 Z_kww = self.Z_dkww[d]
1124 if L > 0:
1125 Ctemp_nw += weight * (
1126 ((Z_knn[k] @ V_knw[k1]) * diagZ_w.conj() +
1127 (dag(Z_knn[k2]) @ V_knw[k2]) * diagZ_w) @ dag(U_ww))
1129 if self.functional == 'var':
1130 # Gradient of the variance term, split in two terms
1131 def variance_term_computer(factor):
1132 result = (
1133 self.nwannier * 2 * weight * (
1134 ((Z_knn[k] @ V_knw[k1]) * factor.conj() +
1135 (dag(Z_knn[k2]) @ V_knw[k2]) * factor) @
1136 dag(U_ww)) / Nw**2
1137 )
1138 return result
1140 first_term = \
1141 O_sum * variance_term_computer(diagZ_w) / Nw**2
1143 second_term = \
1144 - variance_term_computer(diagOZ_w) / Nw
1146 Ctemp_nw += first_term + second_term
1148 temp = Zii_ww.T * Z_kww[k].conj() - Zii_ww * Z_kww[k2].conj()
1149 Utemp_ww += weight * (temp - dag(temp))
1151 if self.functional == 'var':
1152 Utemp_ww += (self.nwannier * 2 * O_sum * weight *
1153 (temp - dag(temp)) / Nw**2)
1155 temp = (OZii_ww.T * Z_kww[k].conj()
1156 - OZii_ww * Z_kww[k2].conj())
1157 Utemp_ww -= (self.nwannier * 2 * weight *
1158 (temp - dag(temp)) / Nw)
1160 dU.append(Utemp_ww.ravel())
1162 if L > 0:
1163 # Ctemp now has same dimension as V, the gradient is in the
1164 # lower-right (Nb-M) x L block
1165 Ctemp_ul = Ctemp_nw[M:, M:]
1166 G_ul = Ctemp_ul - ((C_ul @ dag(C_ul)) @ Ctemp_ul)
1167 dC.append(G_ul.ravel())
1169 return np.concatenate(dU + dC)
1171 def _spread_contributions(self):
1172 """
1173 Compute the contribution of each WF to the spread functional.
1174 """
1175 return (square_modulus_of_Z_diagonal(self.Z_dww).T
1176 @ self.weight_d).real
1178 def step(self, dX, updaterot=True, updatecoeff=True):
1179 # dX is (A, dC) where U->Uexp(-A) and C->C+dC
1180 Nw = self.nwannier
1181 Nk = self.Nk
1182 M_k = self.fixedstates_k
1183 L_k = self.edf_k
1184 if updaterot:
1185 A_kww = dX[:Nk * Nw**2].reshape(Nk, Nw, Nw)
1186 for U, A in zip(self.U_kww, A_kww):
1187 H = -1.j * A.conj()
1188 epsilon, Z = np.linalg.eigh(H)
1189 # Z contains the eigenvectors as COLUMNS.
1190 # Since H = iA, dU = exp(-A) = exp(iH) = ZDZ^d
1191 dU = Z * np.exp(1.j * epsilon) @ dag(Z)
1192 if U.dtype == float:
1193 U[:] = (U @ dU).real
1194 else:
1195 U[:] = U @ dU
1197 if updatecoeff:
1198 start = 0
1199 for C, unocc, L in zip(self.C_kul, self.nbands - M_k, L_k):
1200 if L == 0 or unocc == 0:
1201 continue
1202 Ncoeff = L * unocc
1203 deltaC = dX[Nk * Nw**2 + start: Nk * Nw**2 + start + Ncoeff]
1204 C += deltaC.reshape(unocc, L)
1205 gram_schmidt(C)
1206 start += Ncoeff
1208 self.update()