Coverage for /builds/debichem-team/python-ase/ase/dft/wannier.py: 58.84%
554 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-03-06 04:00 +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 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 _ in range(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(
346 skpoints_kc[n + 1, c] - skpoints_kc[n, c]
347 for n in range(Nk - 1)
348 )
350 for d, Gdir_c in enumerate(Gdir_dc):
351 for k, k_c in enumerate(kpt_kc):
352 # setup dist vector to next kpoint
353 G_c = np.where(Gdir_c > 0, kdist_c, 0)
354 if max(G_c) < 1e-4:
355 kklst_dk[d, k] = k
356 k0_dkc[d, k] = Gdir_c
357 else:
358 kklst_dk[d, k], k0_dkc[d, k] = \
359 neighbor_k_search(k_c, G_c, kpt_kc)
360 return kklst_dk, k0_dkc
363def get_invkklst(kklst_dk):
364 Ndir, Nk = kklst_dk.shape
365 invkklst_dk = np.empty(kklst_dk.shape, int)
366 for d in range(Ndir):
367 for k1 in range(Nk):
368 invkklst_dk[d, k1] = kklst_dk[d].tolist().index(k1)
369 return invkklst_dk
372def choose_states(calcdata, fixedenergy, fixedstates, Nk, nwannier, log, spin):
374 if fixedenergy is None and fixedstates is not None:
375 if isinstance(fixedstates, int):
376 fixedstates = [fixedstates] * Nk
377 fixedstates_k = np.array(fixedstates, int)
378 elif fixedenergy is not None and fixedstates is None:
379 # Setting number of fixed states and EDF from given energy cutoff.
380 # All states below this energy cutoff are fixed.
381 # The reference energy is Ef for metals and CBM for insulators.
382 if calcdata.gap < 0.01 or fixedenergy < 0.01:
383 cutoff = fixedenergy + calcdata.fermi_level
384 else:
385 cutoff = fixedenergy + calcdata.lumo
387 # Find the states below the energy cutoff at each k-point
388 tmp_fixedstates_k = []
389 for k in range(Nk):
390 eps_n = calcdata.eps_skn[spin, k]
391 kindex = eps_n.searchsorted(cutoff)
392 tmp_fixedstates_k.append(kindex)
393 fixedstates_k = np.array(tmp_fixedstates_k, int)
394 elif fixedenergy is not None and fixedstates is not None:
395 raise RuntimeError(
396 'You can not set both fixedenergy and fixedstates')
398 if nwannier == 'auto':
399 if fixedenergy is None and fixedstates is None:
400 # Assume the fixedexergy parameter equal to 0 and
401 # find the states below the Fermi level at each k-point.
402 log("nwannier=auto but no 'fixedenergy' or 'fixedstates'",
403 "parameter was provided, using Fermi level as",
404 "energy cutoff.")
405 tmp_fixedstates_k = []
406 for k in range(Nk):
407 eps_n = calcdata.eps_skn[spin, k]
408 kindex = eps_n.searchsorted(calcdata.fermi_level)
409 tmp_fixedstates_k.append(kindex)
410 fixedstates_k = np.array(tmp_fixedstates_k, int)
411 nwannier = np.max(fixedstates_k)
413 # Without user choice just set nwannier fixed states without EDF
414 if fixedstates is None and fixedenergy is None:
415 fixedstates_k = np.array([nwannier] * Nk, int)
417 return fixedstates_k, nwannier
420def get_eigenvalues(calc):
421 nspins = calc.get_number_of_spins()
422 nkpts = len(calc.get_ibz_k_points())
423 nbands = calc.get_number_of_bands()
424 eps_skn = np.empty((nspins, nkpts, nbands))
426 for ispin in range(nspins):
427 for ikpt in range(nkpts):
428 eps_skn[ispin, ikpt] = calc.get_eigenvalues(kpt=ikpt, spin=ispin)
429 return eps_skn
432class CalcData:
433 def __init__(self, kpt_kc, atoms, fermi_level, lumo, eps_skn,
434 gap):
435 self.kpt_kc = kpt_kc
436 self.atoms = atoms
437 self.fermi_level = fermi_level
438 self.lumo = lumo
439 self.eps_skn = eps_skn
440 self.gap = gap
442 @property
443 def nbands(self):
444 return self.eps_skn.shape[2]
447def get_calcdata(calc):
448 kpt_kc = calc.get_bz_k_points()
449 # Make sure there is no symmetry reduction
450 assert len(calc.get_ibz_k_points()) == len(kpt_kc)
451 lumo = calc.get_homo_lumo()[1]
452 gap = bandgap(calc=calc)[0]
453 return CalcData(
454 kpt_kc=kpt_kc,
455 atoms=calc.get_atoms(),
456 fermi_level=calc.get_fermi_level(),
457 lumo=lumo,
458 eps_skn=get_eigenvalues(calc),
459 gap=gap)
462class Wannier:
463 """Partly occupied Wannier functions
465 Find the set of partly occupied Wannier functions according to
466 Thygesen, Hansen and Jacobsen PRB v72 i12 p125119 2005.
467 """
469 def __init__(self, nwannier, calc,
470 file=None,
471 nbands=None,
472 fixedenergy=None,
473 fixedstates=None,
474 spin=0,
475 initialwannier='orbitals',
476 functional='std',
477 rng=np.random,
478 log=silent):
479 """
480 Required arguments:
482 ``nwannier``: The number of Wannier functions you wish to construct.
483 This must be at least half the number of electrons in the system
484 and at most equal to the number of bands in the calculation.
485 It can also be set to 'auto' in order to automatically choose the
486 minimum number of needed Wannier function based on the
487 ``fixedenergy`` / ``fixedstates`` parameter.
489 ``calc``: A converged DFT calculator class.
490 If ``file`` arg. is not provided, the calculator *must* provide the
491 method ``get_wannier_localization_matrix``, and contain the
492 wavefunctions (save files with only the density is not enough).
493 If the localization matrix is read from file, this is not needed,
494 unless ``get_function`` or ``write_cube`` is called.
496 Optional arguments:
498 ``nbands``: Bands to include in localization.
499 The number of bands considered by Wannier can be smaller than the
500 number of bands in the calculator. This is useful if the highest
501 bands of the DFT calculation are not well converged.
503 ``spin``: The spin channel to be considered.
504 The Wannier code treats each spin channel independently.
506 ``fixedenergy`` / ``fixedstates``: Fixed part of Hilbert space.
507 Determine the fixed part of Hilbert space by either a maximal
508 energy *or* a number of bands (possibly a list for multiple
509 k-points).
510 Default is None meaning that the number of fixed states is equated
511 to ``nwannier``.
512 The maximal energy is relative to the CBM if there is a finite
513 bandgap or to the Fermi level if there is none.
515 ``file``: Read localization and rotation matrices from this file.
517 ``initialwannier``: Initial guess for Wannier rotation matrix.
518 Can be 'bloch' to start from the Bloch states, 'random' to be
519 randomized, 'orbitals' to start from atom-centered d-orbitals and
520 randomly placed gaussian centers (see init_orbitals()),
521 'scdm' to start from localized state selected with SCDM
522 or a list passed to calc.get_initial_wannier.
524 ``functional``: The functional used to measure the localization.
525 Can be 'std' for the standard quadratic functional from the PRB
526 paper, 'var' to add a variance minimizing term.
528 ``rng``: Random number generator for ``initialwannier``.
530 ``log``: Function which logs, such as print().
531 """
532 # Bloch phase sign convention.
533 # May require special cases depending on which code is used.
534 sign = -1
536 self.log = log
537 self.calc = calc
539 self.spin = spin
540 self.functional = functional
541 self.initialwannier = initialwannier
542 self.log('Using functional:', functional)
544 self.calcdata = get_calcdata(calc)
546 self.kptgrid = get_monkhorst_pack_size_and_offset(self.kpt_kc)[0]
547 self.calcdata.kpt_kc *= sign
549 self.largeunitcell_cc = (self.unitcell_cc.T * self.kptgrid).T
550 self.weight_d, self.Gdir_dc = calculate_weights(self.largeunitcell_cc)
551 assert len(self.weight_d) == len(self.Gdir_dc)
553 if nbands is None:
554 # XXX Can work with other number of bands than calculator.
555 # Is this case properly tested, lest we confuse them?
556 nbands = self.calcdata.nbands
557 self.nbands = nbands
559 self.fixedstates_k, self.nwannier = choose_states(
560 self.calcdata, fixedenergy, fixedstates, self.Nk, nwannier,
561 log, spin)
563 # Compute the number of extra degrees of freedom (EDF)
564 self.edf_k = self.nwannier - self.fixedstates_k
566 self.log(f'Wannier: Fixed states : {self.fixedstates_k}')
567 self.log(f'Wannier: Extra degrees of freedom: {self.edf_k}')
569 self.kklst_dk, k0_dkc = get_kklst(self.kpt_kc, self.Gdir_dc)
571 # Set the inverse list of neighboring k-points
572 self.invkklst_dk = get_invkklst(self.kklst_dk)
574 Nw = self.nwannier
575 Nb = self.nbands
576 self.Z_dkww = np.empty((self.Ndir, self.Nk, Nw, Nw), complex)
577 self.V_knw = np.zeros((self.Nk, Nb, Nw), complex)
579 if file is None:
580 self.Z_dknn = self.new_Z(calc, k0_dkc)
581 self.initialize(file=file, initialwannier=initialwannier, rng=rng)
583 @property
584 def atoms(self):
585 return self.calcdata.atoms
587 @property
588 def kpt_kc(self):
589 return self.calcdata.kpt_kc
591 @property
592 def Ndir(self):
593 return len(self.weight_d) # Number of directions
595 @property
596 def Nk(self):
597 return len(self.kpt_kc)
599 def new_Z(self, calc, k0_dkc):
600 Nb = self.nbands
601 Z_dknn = np.empty((self.Ndir, self.Nk, Nb, Nb), complex)
602 for d, dirG in enumerate(self.Gdir_dc):
603 for k in range(self.Nk):
604 k1 = self.kklst_dk[d, k]
605 k0_c = k0_dkc[d, k]
606 Z_dknn[d, k] = calc.get_wannier_localization_matrix(
607 nbands=Nb, dirG=dirG, kpoint=k, nextkpoint=k1,
608 G_I=k0_c, spin=self.spin)
609 return Z_dknn
611 @property
612 def unitcell_cc(self):
613 return self.atoms.cell
615 @property
616 def U_kww(self):
617 return self.wannier_state.U_kww
619 @property
620 def C_kul(self):
621 return self.wannier_state.C_kul
623 def initialize(self, file=None, initialwannier='random', rng=np.random):
624 """Re-initialize current rotation matrix.
626 Keywords are identical to those of the constructor.
627 """
628 from ase.dft.wannierstate import WannierSpec, WannierState
630 spec = WannierSpec(self.Nk, self.nwannier, self.nbands,
631 self.fixedstates_k)
633 if file is not None:
634 with paropen(file, 'r') as fd:
635 Z_dknn, U_kww, C_kul = read_json(fd, always_array=False)
636 self.Z_dknn = Z_dknn
637 wannier_state = WannierState(C_kul, U_kww)
638 elif initialwannier == 'bloch':
639 # Set U and C to pick the lowest Bloch states
640 wannier_state = spec.bloch(self.edf_k)
641 elif initialwannier == 'random':
642 wannier_state = spec.random(rng, self.edf_k)
643 elif initialwannier == 'orbitals':
644 orbitals = init_orbitals(self.atoms, self.nwannier, rng)
645 wannier_state = spec.initial_orbitals(
646 self.calc, orbitals, self.kptgrid, self.edf_k, self.spin)
647 elif initialwannier == 'scdm':
648 wannier_state = spec.scdm(self.calc, self.kpt_kc, self.spin)
649 else:
650 wannier_state = spec.initial_wannier(
651 self.calc, initialwannier, self.kptgrid,
652 self.edf_k, self.spin)
654 self.wannier_state = wannier_state
655 self.update()
657 def save(self, file):
658 """Save information on localization and rotation matrices to file."""
659 with paropen(file, 'w') as fd:
660 write_json(fd, (self.Z_dknn, self.U_kww, self.C_kul))
662 def update(self):
663 # Update large rotation matrix V (from rotation U and coeff C)
664 for k, M in enumerate(self.fixedstates_k):
665 self.V_knw[k, :M] = self.U_kww[k, :M]
666 if M < self.nwannier:
667 self.V_knw[k, M:] = self.C_kul[k] @ self.U_kww[k, M:]
668 # else: self.V_knw[k, M:] = 0.0
670 # Calculate the Zk matrix from the large rotation matrix:
671 # Zk = V^d[k] Zbloch V[k1]
672 for d in range(self.Ndir):
673 for k in range(self.Nk):
674 k1 = self.kklst_dk[d, k]
675 self.Z_dkww[d, k] = dag(self.V_knw[k]) \
676 @ (self.Z_dknn[d, k] @ self.V_knw[k1])
678 # Update the new Z matrix
679 self.Z_dww = self.Z_dkww.sum(axis=1) / self.Nk
681 def get_optimal_nwannier(self, nwrange=5, random_reps=5, tolerance=1e-6):
682 """
683 The optimal value for 'nwannier', maybe.
685 The optimal value is the one that gives the lowest average value for
686 the spread of the most delocalized Wannier function in the set.
688 ``nwrange``: number of different values to try for 'nwannier', the
689 values will span a symmetric range around ``nwannier`` if possible.
691 ``random_reps``: number of repetitions with random seed, the value is
692 then an average over these repetitions.
694 ``tolerance``: tolerance for the gradient descent algorithm, can be
695 useful to increase the speed, with a cost in accuracy.
696 """
698 # Define the range of values to try based on the maximum number of fixed
699 # states (that is the minimum number of WFs we need) and the number of
700 # available bands we have.
701 max_number_fixedstates = np.max(self.fixedstates_k)
703 min_range_value = max(self.nwannier - int(np.floor(nwrange / 2)),
704 max_number_fixedstates)
705 max_range_value = min(min_range_value + nwrange, self.nbands + 1)
706 Nws = np.arange(min_range_value, max_range_value)
708 # If there is no randomness, there is no need to repeat
709 random_initials = ['random', 'orbitals']
710 if self.initialwannier not in random_initials:
711 random_reps = 1
713 t = -time()
714 avg_max_spreads = np.zeros(len(Nws))
715 for j, Nw in enumerate(Nws):
716 self.log('Trying with Nw =', Nw)
718 # Define once with the fastest 'initialwannier',
719 # then initialize with random seeds in the for loop
720 wan = Wannier(nwannier=int(Nw),
721 calc=self.calc,
722 nbands=self.nbands,
723 spin=self.spin,
724 functional=self.functional,
725 initialwannier='bloch',
726 log=self.log,
727 rng=self.rng)
728 wan.fixedstates_k = self.fixedstates_k
729 wan.edf_k = wan.nwannier - wan.fixedstates_k
731 max_spreads = np.zeros(random_reps)
732 for i in range(random_reps):
733 wan.initialize(initialwannier=self.initialwannier)
734 wan.localize(tolerance=tolerance)
735 max_spreads[i] = np.max(wan.get_spreads())
737 avg_max_spreads[j] = max_spreads.mean()
739 self.log('Average spreads: ', avg_max_spreads)
740 t += time()
741 self.log(f'Execution time: {t:.1f}s')
743 return Nws[np.argmin(avg_max_spreads)]
745 def get_centers(self, scaled=False):
746 """Calculate the Wannier centers
748 ::
750 pos = L / 2pi * phase(diag(Z))
751 """
752 coord_wc = np.angle(self.Z_dww[:3].diagonal(0, 1, 2)).T / (2 * pi) % 1
753 if not scaled:
754 coord_wc = coord_wc @ self.largeunitcell_cc
755 return coord_wc
757 def get_radii(self):
758 r"""Calculate the spread of the Wannier functions.
760 ::
762 -- / L \ 2 2
763 radius**2 = - > | --- | ln |Z|
764 --d \ 2pi /
766 Note that this function can fail with some Bravais lattices,
767 see `get_spreads()` for a more robust alternative.
768 """
769 r2 = - (self.largeunitcell_cc.diagonal()**2 / (2 * pi)**2) \
770 @ np.log(abs(self.Z_dww[:3].diagonal(0, 1, 2))**2)
771 return np.sqrt(r2)
773 def get_spreads(self):
774 r"""Calculate the spread of the Wannier functions in Ų.
775 The definition is based on eq. 13 in PRB61-15 by Berghold and Mundy.
777 ::
779 / 1 \ 2 -- 2
780 spread = - |----| > W_d ln |Z_dw|
781 \2 pi/ --d
784 """
785 # compute weights without normalization, to keep physical dimension
786 weight_d, _ = calculate_weights(self.largeunitcell_cc, normalize=False)
787 Z2_dw = square_modulus_of_Z_diagonal(self.Z_dww)
788 spread_w = - (np.log(Z2_dw).T @ weight_d).real / (2 * np.pi)**2
789 return spread_w
791 def get_spectral_weight(self, w):
792 return abs(self.V_knw[:, :, w])**2 / self.Nk
794 def get_pdos(self, w, energies, width):
795 """Projected density of states (PDOS).
797 Returns the (PDOS) for Wannier function ``w``. The calculation
798 is performed over the energy grid specified in energies. The
799 PDOS is produced as a sum of Gaussians centered at the points
800 of the energy grid and with the specified width.
801 """
802 spec_kn = self.get_spectral_weight(w)
803 dos = np.zeros(len(energies))
804 for k, spec_n in enumerate(spec_kn):
805 eig_n = self.calcdata.eps_skn[self.spin, k]
806 for weight, eig in zip(spec_n, eig_n):
807 # Add gaussian centered at the eigenvalue
808 x = ((energies - eig) / width)**2
809 dos += weight * np.exp(-x.clip(0., 40.)) / (sqrt(pi) * width)
810 return dos
812 def translate(self, w, R):
813 """Translate the w'th Wannier function
815 The distance vector R = [n1, n2, n3], is in units of the basis
816 vectors of the small cell.
817 """
818 for kpt_c, U_ww in zip(self.kpt_kc, self.U_kww):
819 U_ww[:, w] *= np.exp(2.j * pi * (np.array(R) @ kpt_c))
820 self.update()
822 def translate_to_cell(self, w, cell):
823 """Translate the w'th Wannier function to specified cell"""
824 scaled_c = np.angle(self.Z_dww[:3, w, w]) * self.kptgrid / (2 * pi)
825 trans = np.array(cell) - np.floor(scaled_c)
826 self.translate(w, trans)
828 def translate_all_to_cell(self, cell=(0, 0, 0)):
829 r"""Translate all Wannier functions to specified cell.
831 Move all Wannier orbitals to a specific unit cell. There
832 exists an arbitrariness in the positions of the Wannier
833 orbitals relative to the unit cell. This method can move all
834 orbitals to the unit cell specified by ``cell``. For a
835 `\Gamma`-point calculation, this has no effect. For a
836 **k**-point calculation the periodicity of the orbitals are
837 given by the large unit cell defined by repeating the original
838 unitcell by the number of **k**-points in each direction. In
839 this case it is useful to move the orbitals away from the
840 boundaries of the large cell before plotting them. For a bulk
841 calculation with, say 10x10x10 **k** points, one could move
842 the orbitals to the cell [2,2,2]. In this way the pbc
843 boundary conditions will not be noticed.
844 """
845 scaled_wc = (np.angle(self.Z_dww[:3].diagonal(0, 1, 2)).T *
846 self.kptgrid / (2 * pi))
847 trans_wc = np.array(cell)[None] - np.floor(scaled_wc)
848 for kpt_c, U_ww in zip(self.kpt_kc, self.U_kww):
849 U_ww *= np.exp(2.j * pi * (trans_wc @ kpt_c))
850 self.update()
852 def distances(self, R):
853 """Relative distances between centers.
855 Returns a matrix with the distances between different Wannier centers.
856 R = [n1, n2, n3] is in units of the basis vectors of the small cell
857 and allows one to measure the distance with centers moved to a
858 different small cell.
859 The dimension of the matrix is [Nw, Nw].
860 """
861 Nw = self.nwannier
862 cen = self.get_centers()
863 r1 = cen.repeat(Nw, axis=0).reshape(Nw, Nw, 3)
864 r2 = cen.copy()
865 for i in range(3):
866 r2 += self.unitcell_cc[i] * R[i]
868 r2 = np.swapaxes(r2.repeat(Nw, axis=0).reshape(Nw, Nw, 3), 0, 1)
869 return np.sqrt(np.sum((r1 - r2)**2, axis=-1))
871 @functools.lru_cache(maxsize=10000)
872 def _get_hopping(self, n1, n2, n3):
873 """Returns the matrix H(R)_nm=<0,n|H|R,m>.
875 ::
877 1 _ -ik.R
878 H(R) = <0,n|H|R,m> = --- >_ e H(k)
879 Nk k
881 where R = (n1, n2, n3) is the cell-distance (in units of the basis
882 vectors of the small cell) and n,m are indices of the Wannier functions.
883 This function caches up to 'maxsize' results.
884 """
885 R = np.array([n1, n2, n3], float)
886 H_ww = np.zeros([self.nwannier, self.nwannier], complex)
887 for k, kpt_c in enumerate(self.kpt_kc):
888 phase = np.exp(-2.j * pi * (np.array(R) @ kpt_c))
889 H_ww += self.get_hamiltonian(k) * phase
890 return H_ww / self.Nk
892 def get_hopping(self, R):
893 """Returns the matrix H(R)_nm=<0,n|H|R,m>.
895 ::
897 1 _ -ik.R
898 H(R) = <0,n|H|R,m> = --- >_ e H(k)
899 Nk k
901 where R is the cell-distance (in units of the basis vectors of
902 the small cell) and n,m are indices of the Wannier functions.
903 """
904 return self._get_hopping(R[0], R[1], R[2])
906 def get_hamiltonian(self, k):
907 """Get Hamiltonian at existing k-vector of index k
909 ::
911 dag
912 H(k) = V diag(eps ) V
913 k k k
914 """
915 eps_n = self.calcdata.eps_skn[self.spin, k, :self.nbands]
916 V_nw = self.V_knw[k]
917 return (dag(V_nw) * eps_n) @ V_nw
919 def get_hamiltonian_kpoint(self, kpt_c):
920 """Get Hamiltonian at some new arbitrary k-vector
922 ::
924 _ ik.R
925 H(k) = >_ e H(R)
926 R
928 Warning: This method moves all Wannier functions to cell (0, 0, 0)
929 """
930 self.log('Translating all Wannier functions to cell (0, 0, 0)')
931 self.translate_all_to_cell()
932 max = (self.kptgrid - 1) // 2
933 N1, N2, N3 = max
934 Hk = np.zeros([self.nwannier, self.nwannier], complex)
935 for n1 in range(-N1, N1 + 1):
936 for n2 in range(-N2, N2 + 1):
937 for n3 in range(-N3, N3 + 1):
938 R = np.array([n1, n2, n3], float)
939 hop_ww = self.get_hopping(R)
940 phase = np.exp(+2.j * pi * (R @ kpt_c))
941 Hk += hop_ww * phase
942 return Hk
944 def get_function(self, index, repeat=None):
945 r"""Get Wannier function on grid.
947 Returns an array with the funcion values of the indicated Wannier
948 function on a grid with the size of the *repeated* unit cell.
950 For a calculation using **k**-points the relevant unit cell for
951 eg. visualization of the Wannier orbitals is not the original unit
952 cell, but rather a larger unit cell defined by repeating the
953 original unit cell by the number of **k**-points in each direction.
954 Note that for a `\Gamma`-point calculation the large unit cell
955 coinsides with the original unit cell.
956 The large unitcell also defines the periodicity of the Wannier
957 orbitals.
959 ``index`` can be either a single WF or a coordinate vector in terms
960 of the WFs.
961 """
963 # Default size of plotting cell is the one corresponding to k-points.
964 if repeat is None:
965 repeat = self.kptgrid
966 N1, N2, N3 = repeat
968 dim = self.calc.get_number_of_grid_points()
969 largedim = dim * [N1, N2, N3]
971 wanniergrid = np.zeros(largedim, dtype=complex)
972 for k, kpt_c in enumerate(self.kpt_kc):
973 # The coordinate vector of wannier functions
974 if isinstance(index, int):
975 vec_n = self.V_knw[k, :, index]
976 else:
977 vec_n = self.V_knw[k] @ index
979 wan_G = np.zeros(dim, complex)
980 for n, coeff in enumerate(vec_n):
981 wan_G += coeff * self.calc.get_pseudo_wave_function(
982 n, k, self.spin, pad=True)
984 # Distribute the small wavefunction over large cell:
985 for n1 in range(N1):
986 for n2 in range(N2):
987 for n3 in range(N3): # sign?
988 e = np.exp(-2.j * pi * np.array([n1, n2, n3]) @ kpt_c)
989 wanniergrid[n1 * dim[0]:(n1 + 1) * dim[0],
990 n2 * dim[1]:(n2 + 1) * dim[1],
991 n3 * dim[2]:(n3 + 1) * dim[2]] += e * wan_G
993 # Normalization
994 wanniergrid /= np.sqrt(self.Nk)
995 return wanniergrid
997 def write_cube(self, index, fname, repeat=None, angle=False):
998 """
999 Dump specified Wannier function to a cube file.
1001 Arguments:
1003 ``index``: Integer, index of the Wannier function to save.
1005 ``repeat``: Array of integer, repeat supercell and Wannier function.
1007 ``fname``: Name of the cube file.
1009 ``angle``: If False, save the absolute value. If True, save
1010 the complex phase of the Wannier function.
1011 """
1012 from ase.io import write
1014 # Default size of plotting cell is the one corresponding to k-points.
1015 if repeat is None:
1016 repeat = self.kptgrid
1018 # Remove constraints, some are not compatible with repeat()
1019 atoms = self.atoms.copy()
1020 atoms.set_constraint()
1021 atoms = atoms * repeat
1022 func = self.get_function(index, repeat)
1024 # Compute absolute value or complex angle
1025 if angle:
1026 data = np.angle(func)
1027 else:
1028 if self.Nk == 1:
1029 func *= np.exp(-1.j * np.angle(func.max()))
1030 func = abs(func)
1031 data = func
1033 write(fname, atoms, data=data, format='cube')
1035 def localize(self, step=0.25, tolerance=1e-08,
1036 updaterot=True, updatecoeff=True):
1037 """Optimize rotation to give maximal localization"""
1038 md_min(self, step=step, tolerance=tolerance, log=self.log,
1039 updaterot=updaterot, updatecoeff=updatecoeff)
1041 def get_functional_value(self):
1042 """Calculate the value of the spread functional.
1044 ::
1046 Tr[|ZI|^2]=sum(I)sum(n) w_i|Z_(i)_nn|^2,
1048 where w_i are weights.
1050 If the functional is set to 'var' it subtracts a variance term
1052 ::
1054 Nw * var(sum(n) w_i|Z_(i)_nn|^2),
1056 where Nw is the number of WFs ``nwannier``.
1057 """
1058 a_w = self._spread_contributions()
1059 if self.functional == 'std':
1060 fun = np.sum(a_w)
1061 elif self.functional == 'var':
1062 fun = np.sum(a_w) - self.nwannier * np.var(a_w)
1063 self.log(f'std: {np.sum(a_w):.4f}',
1064 f'\tvar: {self.nwannier * np.var(a_w):.4f}')
1065 return fun
1067 def get_gradients(self):
1068 # Determine gradient of the spread functional.
1069 #
1070 # The gradient for a rotation A_kij is::
1071 #
1072 # dU = dRho/dA_{k,i,j} = sum(I) sum(k')
1073 # + Z_jj Z_kk',ij^* - Z_ii Z_k'k,ij^*
1074 # - Z_ii^* Z_kk',ji + Z_jj^* Z_k'k,ji
1075 #
1076 # The gradient for a change of coefficients is::
1077 #
1078 # dRho/da^*_{k,i,j} = sum(I) [[(Z_0)_{k} V_{k'} diag(Z^*) +
1079 # (Z_0_{k''})^d V_{k''} diag(Z)] *
1080 # U_k^d]_{N+i,N+j}
1081 #
1082 # where diag(Z) is a square,diagonal matrix with Z_nn in the diagonal,
1083 # k' = k + dk and k = k'' + dk.
1084 #
1085 # The extra degrees of freedom chould be kept orthonormal to the fixed
1086 # space, thus we introduce lagrange multipliers, and minimize instead::
1087 #
1088 # Rho_L = Rho - sum_{k,n,m} lambda_{k,nm} <c_{kn}|c_{km}>
1089 #
1090 # for this reason the coefficient gradients should be multiplied
1091 # by (1 - c c^dag).
1093 Nb = self.nbands
1094 Nw = self.nwannier
1096 if self.functional == 'var':
1097 O_w = self._spread_contributions()
1098 O_sum = np.sum(O_w)
1100 dU = []
1101 dC = []
1102 for k in range(self.Nk):
1103 M = self.fixedstates_k[k]
1104 L = self.edf_k[k]
1105 U_ww = self.U_kww[k]
1106 C_ul = self.C_kul[k]
1107 Utemp_ww = np.zeros((Nw, Nw), complex)
1108 Ctemp_nw = np.zeros((Nb, Nw), complex)
1110 for d, weight in enumerate(self.weight_d):
1111 if abs(weight) < 1.0e-6:
1112 continue
1114 Z_knn = self.Z_dknn[d]
1115 diagZ_w = self.Z_dww[d].diagonal()
1116 Zii_ww = np.repeat(diagZ_w, Nw).reshape(Nw, Nw)
1117 if self.functional == 'var':
1118 diagOZ_w = O_w * diagZ_w
1119 OZii_ww = np.repeat(diagOZ_w, Nw).reshape(Nw, Nw)
1121 k1 = self.kklst_dk[d, k]
1122 k2 = self.invkklst_dk[d, k]
1123 V_knw = self.V_knw
1124 Z_kww = self.Z_dkww[d]
1126 if L > 0:
1127 Ctemp_nw += weight * (
1128 ((Z_knn[k] @ V_knw[k1]) * diagZ_w.conj() +
1129 (dag(Z_knn[k2]) @ V_knw[k2]) * diagZ_w) @ dag(U_ww))
1131 if self.functional == 'var':
1132 # Gradient of the variance term, split in two terms
1133 def variance_term_computer(factor):
1134 result = (
1135 self.nwannier * 2 * weight * (
1136 ((Z_knn[k] @ V_knw[k1]) * factor.conj() +
1137 (dag(Z_knn[k2]) @ V_knw[k2]) * factor) @
1138 dag(U_ww)) / Nw**2
1139 )
1140 return result
1142 first_term = \
1143 O_sum * variance_term_computer(diagZ_w) / Nw**2
1145 second_term = \
1146 - variance_term_computer(diagOZ_w) / Nw
1148 Ctemp_nw += first_term + second_term
1150 temp = Zii_ww.T * Z_kww[k].conj() - Zii_ww * Z_kww[k2].conj()
1151 Utemp_ww += weight * (temp - dag(temp))
1153 if self.functional == 'var':
1154 Utemp_ww += (self.nwannier * 2 * O_sum * weight *
1155 (temp - dag(temp)) / Nw**2)
1157 temp = (OZii_ww.T * Z_kww[k].conj()
1158 - OZii_ww * Z_kww[k2].conj())
1159 Utemp_ww -= (self.nwannier * 2 * weight *
1160 (temp - dag(temp)) / Nw)
1162 dU.append(Utemp_ww.ravel())
1164 if L > 0:
1165 # Ctemp now has same dimension as V, the gradient is in the
1166 # lower-right (Nb-M) x L block
1167 Ctemp_ul = Ctemp_nw[M:, M:]
1168 G_ul = Ctemp_ul - ((C_ul @ dag(C_ul)) @ Ctemp_ul)
1169 dC.append(G_ul.ravel())
1171 return np.concatenate(dU + dC)
1173 def _spread_contributions(self):
1174 """
1175 Compute the contribution of each WF to the spread functional.
1176 """
1177 return (square_modulus_of_Z_diagonal(self.Z_dww).T
1178 @ self.weight_d).real
1180 def step(self, dX, updaterot=True, updatecoeff=True):
1181 # dX is (A, dC) where U->Uexp(-A) and C->C+dC
1182 Nw = self.nwannier
1183 Nk = self.Nk
1184 M_k = self.fixedstates_k
1185 L_k = self.edf_k
1186 if updaterot:
1187 A_kww = dX[:Nk * Nw**2].reshape(Nk, Nw, Nw)
1188 for U, A in zip(self.U_kww, A_kww):
1189 H = -1.j * A.conj()
1190 epsilon, Z = np.linalg.eigh(H)
1191 # Z contains the eigenvectors as COLUMNS.
1192 # Since H = iA, dU = exp(-A) = exp(iH) = ZDZ^d
1193 dU = Z * np.exp(1.j * epsilon) @ dag(Z)
1194 if U.dtype == float:
1195 U[:] = (U @ dU).real
1196 else:
1197 U[:] = U @ dU
1199 if updatecoeff:
1200 start = 0
1201 for C, unocc, L in zip(self.C_kul, self.nbands - M_k, L_k):
1202 if L == 0 or unocc == 0:
1203 continue
1204 Ncoeff = L * unocc
1205 deltaC = dX[Nk * Nw**2 + start: Nk * Nw**2 + start + Ncoeff]
1206 C += deltaC.reshape(unocc, L)
1207 gram_schmidt(C)
1208 start += Ncoeff
1210 self.update()