Coverage for /builds/debichem-team/python-ase/ase/dft/wannier.py : 31.65%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1""" Maximally localized Wannier Functions
3 Find the set of maximally localized Wannier functions
4 using the spread functional of Marzari and Vanderbilt
5 (PRB 56, 1997 page 12847).
6"""
7from time import time
8from math import sqrt, pi
10import numpy as np
12from ase.parallel import paropen
13from ase.dft.kpoints import get_monkhorst_pack_size_and_offset
14from ase.transport.tools import dagger, normalize
15from ase.io.jsonio import read_json, write_json
17dag = dagger
20def gram_schmidt(U):
21 """Orthonormalize columns of U according to the Gram-Schmidt procedure."""
22 for i, col in enumerate(U.T):
23 for col2 in U.T[:i]:
24 col -= col2 * np.dot(col2.conj(), col)
25 col /= np.linalg.norm(col)
28def lowdin(U, S=None):
29 """Orthonormalize columns of U according to the Lowdin procedure.
31 If the overlap matrix is know, it can be specified in S.
32 """
33 if S is None:
34 S = np.dot(dag(U), U)
35 eig, rot = np.linalg.eigh(S)
36 rot = np.dot(rot / np.sqrt(eig), dag(rot))
37 U[:] = np.dot(U, rot)
40def neighbor_k_search(k_c, G_c, kpt_kc, tol=1e-4):
41 # search for k1 (in kpt_kc) and k0 (in alldir), such that
42 # k1 - k - G + k0 = 0
43 alldir_dc = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1],
44 [1, 1, 0], [1, 0, 1], [0, 1, 1]], dtype=int)
45 for k0_c in alldir_dc:
46 for k1, k1_c in enumerate(kpt_kc):
47 if np.linalg.norm(k1_c - k_c - G_c + k0_c) < tol:
48 return k1, k0_c
50 print('Wannier: Did not find matching kpoint for kpt=', k_c)
51 print('Probably non-uniform k-point grid')
52 raise NotImplementedError
55def calculate_weights(cell_cc, normalize=True):
56 """ Weights are used for non-cubic cells, see PRB **61**, 10040"""
57 alldirs_dc = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1],
58 [1, 1, 0], [1, 0, 1], [0, 1, 1]], dtype=int)
59 g = np.dot(cell_cc, cell_cc.T)
60 # NOTE: Only first 3 of following 6 weights are presently used:
61 w = np.zeros(6)
62 w[0] = g[0, 0] - g[0, 1] - g[0, 2]
63 w[1] = g[1, 1] - g[0, 1] - g[1, 2]
64 w[2] = g[2, 2] - g[0, 2] - g[1, 2]
65 w[3] = g[0, 1]
66 w[4] = g[0, 2]
67 w[5] = g[1, 2]
68 # Make sure that first 3 Gdir vectors are included -
69 # these are used to calculate Wanniercenters.
70 Gdir_dc = alldirs_dc[:3]
71 weight_d = w[:3]
72 for d in range(3, 6):
73 if abs(w[d]) > 1e-5:
74 Gdir_dc = np.concatenate((Gdir_dc, alldirs_dc[d:d + 1]))
75 weight_d = np.concatenate((weight_d, w[d:d + 1]))
76 if normalize:
77 weight_d /= max(abs(weight_d))
78 return weight_d, Gdir_dc
81def random_orthogonal_matrix(dim, rng=np.random, real=False):
82 """Generate a random orthogonal matrix"""
84 H = rng.rand(dim, dim)
85 np.add(dag(H), H, H)
86 np.multiply(.5, H, H)
88 if real:
89 gram_schmidt(H)
90 return H
91 else:
92 val, vec = np.linalg.eig(H)
93 return np.dot(vec * np.exp(1.j * val), dag(vec))
96def steepest_descent(func, step=.005, tolerance=1e-6, verbose=False, **kwargs):
97 fvalueold = 0.
98 fvalue = fvalueold + 10
99 count = 0
100 while abs((fvalue - fvalueold) / fvalue) > tolerance:
101 fvalueold = fvalue
102 dF = func.get_gradients()
103 func.step(dF * step, **kwargs)
104 fvalue = func.get_functional_value()
105 count += 1
106 if verbose:
107 print('SteepestDescent: iter=%s, value=%s' % (count, fvalue))
110def md_min(func, step=.25, tolerance=1e-6, verbose=False, **kwargs):
111 if verbose:
112 print('Localize with step =', step, 'and tolerance =', tolerance)
113 t = -time()
114 fvalueold = 0.
115 fvalue = fvalueold + 10
116 count = 0
117 V = np.zeros(func.get_gradients().shape, dtype=complex)
118 while abs((fvalue - fvalueold) / fvalue) > tolerance:
119 fvalueold = fvalue
120 dF = func.get_gradients()
121 V *= (dF * V.conj()).real > 0
122 V += step * dF
123 func.step(V, **kwargs)
124 fvalue = func.get_functional_value()
125 if fvalue < fvalueold:
126 step *= 0.5
127 count += 1
128 if verbose:
129 print('MDmin: iter=%s, step=%s, value=%s' % (count, step, fvalue))
130 if verbose:
131 t += time()
132 print('%d iterations in %0.2f seconds (%0.2f ms/iter), endstep = %s' % (
133 count, t, t * 1000. / count, step))
136def rotation_from_projection(proj_nw, fixed, ortho=True):
137 """Determine rotation and coefficient matrices from projections
139 proj_nw = <psi_n|p_w>
140 psi_n: eigenstates
141 p_w: localized function
143 Nb (n) = Number of bands
144 Nw (w) = Number of wannier functions
145 M (f) = Number of fixed states
146 L (l) = Number of extra degrees of freedom
147 U (u) = Number of non-fixed states
148 """
150 Nb, Nw = proj_nw.shape
151 M = fixed
152 L = Nw - M
154 U_ww = np.empty((Nw, Nw), dtype=proj_nw.dtype)
155 U_ww[:M] = proj_nw[:M]
157 if L > 0:
158 proj_uw = proj_nw[M:]
159 eig_w, C_ww = np.linalg.eigh(np.dot(dag(proj_uw), proj_uw))
160 C_ul = np.dot(proj_uw, C_ww[:, np.argsort(-eig_w.real)[:L]])
161 # eig_u, C_uu = np.linalg.eigh(np.dot(proj_uw, dag(proj_uw)))
162 # C_ul = C_uu[:, np.argsort(-eig_u.real)[:L]]
164 U_ww[M:] = np.dot(dag(C_ul), proj_uw)
165 else:
166 C_ul = np.empty((Nb - M, 0))
168 normalize(C_ul)
169 if ortho:
170 lowdin(U_ww)
171 else:
172 normalize(U_ww)
174 return U_ww, C_ul
177class Wannier:
178 """Maximally localized Wannier Functions
180 Find the set of maximally localized Wannier functions using the
181 spread functional of Marzari and Vanderbilt (PRB 56, 1997 page
182 12847).
183 """
185 def __init__(self, nwannier, calc,
186 file=None,
187 nbands=None,
188 fixedenergy=None,
189 fixedstates=None,
190 spin=0,
191 initialwannier='random',
192 rng=np.random,
193 verbose=False):
194 """
195 Required arguments:
197 ``nwannier``: The number of Wannier functions you wish to construct.
198 This must be at least half the number of electrons in the system
199 and at most equal to the number of bands in the calculation.
201 ``calc``: A converged DFT calculator class.
202 If ``file`` arg. is not provided, the calculator *must* provide the
203 method ``get_wannier_localization_matrix``, and contain the
204 wavefunctions (save files with only the density is not enough).
205 If the localization matrix is read from file, this is not needed,
206 unless ``get_function`` or ``write_cube`` is called.
208 Optional arguments:
210 ``nbands``: Bands to include in localization.
211 The number of bands considered by Wannier can be smaller than the
212 number of bands in the calculator. This is useful if the highest
213 bands of the DFT calculation are not well converged.
215 ``spin``: The spin channel to be considered.
216 The Wannier code treats each spin channel independently.
218 ``fixedenergy`` / ``fixedstates``: Fixed part of Heilbert space.
219 Determine the fixed part of Hilbert space by either a maximal
220 energy *or* a number of bands (possibly a list for multiple
221 k-points).
222 Default is None meaning that the number of fixed states is equated
223 to ``nwannier``.
225 ``file``: Read localization and rotation matrices from this file.
227 ``initialwannier``: Initial guess for Wannier rotation matrix.
228 Can be 'bloch' to start from the Bloch states, 'random' to be
229 randomized, or a list passed to calc.get_initial_wannier.
231 ``rng``: Random number generator for ``initialwannier``.
233 ``verbose``: True / False level of verbosity.
234 """
235 # Bloch phase sign convention.
236 # May require special cases depending on which code is used.
237 sign = -1
239 self.nwannier = nwannier
240 self.calc = calc
241 self.spin = spin
242 self.verbose = verbose
243 self.kpt_kc = calc.get_bz_k_points()
244 assert len(calc.get_ibz_k_points()) == len(self.kpt_kc)
245 self.kptgrid = get_monkhorst_pack_size_and_offset(self.kpt_kc)[0]
246 self.kpt_kc *= sign
248 self.Nk = len(self.kpt_kc)
249 self.unitcell_cc = calc.get_atoms().get_cell()
250 self.largeunitcell_cc = (self.unitcell_cc.T * self.kptgrid).T
251 self.weight_d, self.Gdir_dc = calculate_weights(self.largeunitcell_cc)
252 self.Ndir = len(self.weight_d) # Number of directions
254 if nbands is not None:
255 self.nbands = nbands
256 else:
257 self.nbands = calc.get_number_of_bands()
258 if fixedenergy is None:
259 if fixedstates is None:
260 self.fixedstates_k = np.array([nwannier] * self.Nk, int)
261 else:
262 if isinstance(fixedstates, int):
263 fixedstates = [fixedstates] * self.Nk
264 self.fixedstates_k = np.array(fixedstates, int)
265 else:
266 # Setting number of fixed states and EDF from specified energy.
267 # All states below this energy (relative to Fermi level) are fixed.
268 fixedenergy += calc.get_fermi_level()
269 print(fixedenergy)
270 self.fixedstates_k = np.array(
271 [calc.get_eigenvalues(k, spin).searchsorted(fixedenergy)
272 for k in range(self.Nk)], int)
273 self.edf_k = self.nwannier - self.fixedstates_k
274 if verbose:
275 print('Wannier: Fixed states : %s' % self.fixedstates_k)
276 print('Wannier: Extra degrees of freedom: %s' % self.edf_k)
278 # Set the list of neighboring k-points k1, and the "wrapping" k0,
279 # such that k1 - k - G + k0 = 0
280 #
281 # Example: kpoints = (-0.375,-0.125,0.125,0.375), dir=0
282 # G = [0.25,0,0]
283 # k=0.375, k1= -0.375 : -0.375-0.375-0.25 => k0=[1,0,0]
284 #
285 # For a gamma point calculation k1 = k = 0, k0 = [1,0,0] for dir=0
286 if self.Nk == 1:
287 self.kklst_dk = np.zeros((self.Ndir, 1), int)
288 k0_dkc = self.Gdir_dc.reshape(-1, 1, 3)
289 else:
290 self.kklst_dk = np.empty((self.Ndir, self.Nk), int)
291 k0_dkc = np.empty((self.Ndir, self.Nk, 3), int)
293 # Distance between kpoints
294 kdist_c = np.empty(3)
295 for c in range(3):
296 # make a sorted list of the kpoint values in this direction
297 slist = np.argsort(self.kpt_kc[:, c], kind='mergesort')
298 skpoints_kc = np.take(self.kpt_kc, slist, axis=0)
299 kdist_c[c] = max([skpoints_kc[n + 1, c] - skpoints_kc[n, c]
300 for n in range(self.Nk - 1)])
302 for d, Gdir_c in enumerate(self.Gdir_dc):
303 for k, k_c in enumerate(self.kpt_kc):
304 # setup dist vector to next kpoint
305 G_c = np.where(Gdir_c > 0, kdist_c, 0)
306 if max(G_c) < 1e-4:
307 self.kklst_dk[d, k] = k
308 k0_dkc[d, k] = Gdir_c
309 else:
310 self.kklst_dk[d, k], k0_dkc[d, k] = \
311 neighbor_k_search(k_c, G_c, self.kpt_kc)
313 # Set the inverse list of neighboring k-points
314 self.invkklst_dk = np.empty((self.Ndir, self.Nk), int)
315 for d in range(self.Ndir):
316 for k1 in range(self.Nk):
317 self.invkklst_dk[d, k1] = self.kklst_dk[d].tolist().index(k1)
319 Nw = self.nwannier
320 Nb = self.nbands
321 self.Z_dkww = np.empty((self.Ndir, self.Nk, Nw, Nw), complex)
322 self.V_knw = np.zeros((self.Nk, Nb, Nw), complex)
323 if file is None:
324 self.Z_dknn = np.empty((self.Ndir, self.Nk, Nb, Nb), complex)
325 for d, dirG in enumerate(self.Gdir_dc):
326 for k in range(self.Nk):
327 k1 = self.kklst_dk[d, k]
328 k0_c = k0_dkc[d, k]
329 self.Z_dknn[d, k] = calc.get_wannier_localization_matrix(
330 nbands=Nb, dirG=dirG, kpoint=k, nextkpoint=k1,
331 G_I=k0_c, spin=self.spin)
332 self.initialize(file=file, initialwannier=initialwannier, rng=rng)
334 def initialize(self, file=None, initialwannier='random', rng=np.random):
335 """Re-initialize current rotation matrix.
337 Keywords are identical to those of the constructor.
338 """
339 Nw = self.nwannier
340 Nb = self.nbands
342 if file is not None:
343 with paropen(file, 'r') as fd:
344 self.Z_dknn, self.U_kww, self.C_kul = read_json(fd)
346 elif initialwannier == 'bloch':
347 # Set U and C to pick the lowest Bloch states
348 self.U_kww = np.zeros((self.Nk, Nw, Nw), complex)
349 self.C_kul = []
350 for U, M, L in zip(self.U_kww, self.fixedstates_k, self.edf_k):
351 U[:] = np.identity(Nw, complex)
352 if L > 0:
353 self.C_kul.append(
354 np.identity(Nb - M, complex)[:, :L])
355 else:
356 self.C_kul.append([])
357 elif initialwannier == 'random':
358 # Set U and C to random (orthogonal) matrices
359 self.U_kww = np.zeros((self.Nk, Nw, Nw), complex)
360 self.C_kul = []
361 for U, M, L in zip(self.U_kww, self.fixedstates_k, self.edf_k):
362 U[:] = random_orthogonal_matrix(Nw, rng, real=False)
363 if L > 0:
364 self.C_kul.append(random_orthogonal_matrix(
365 Nb - M, rng=rng, real=False)[:, :L])
366 else:
367 self.C_kul.append(np.array([]))
368 else:
369 # Use initial guess to determine U and C
370 self.C_kul, self.U_kww = self.calc.initial_wannier(
371 initialwannier, self.kptgrid, self.fixedstates_k,
372 self.edf_k, self.spin, self.nbands)
373 self.update()
375 def save(self, file):
376 """Save information on localization and rotation matrices to file."""
377 with paropen(file, 'w') as fd:
378 write_json(fd, (self.Z_dknn, self.U_kww, self.C_kul))
380 def update(self):
381 # Update large rotation matrix V (from rotation U and coeff C)
382 for k, M in enumerate(self.fixedstates_k):
383 self.V_knw[k, :M] = self.U_kww[k, :M]
384 if M < self.nwannier:
385 self.V_knw[k, M:] = np.dot(self.C_kul[k], self.U_kww[k, M:])
386 # else: self.V_knw[k, M:] = 0.0
388 # Calculate the Zk matrix from the large rotation matrix:
389 # Zk = V^d[k] Zbloch V[k1]
390 for d in range(self.Ndir):
391 for k in range(self.Nk):
392 k1 = self.kklst_dk[d, k]
393 self.Z_dkww[d, k] = np.dot(dag(self.V_knw[k]), np.dot(
394 self.Z_dknn[d, k], self.V_knw[k1]))
396 # Update the new Z matrix
397 self.Z_dww = self.Z_dkww.sum(axis=1) / self.Nk
399 def get_centers(self, scaled=False):
400 """Calculate the Wannier centers
402 ::
404 pos = L / 2pi * phase(diag(Z))
405 """
406 coord_wc = np.angle(self.Z_dww[:3].diagonal(0, 1, 2)).T / (2 * pi) % 1
407 if not scaled:
408 coord_wc = np.dot(coord_wc, self.largeunitcell_cc)
409 return coord_wc
411 def get_radii(self):
412 r"""Calculate the spread of the Wannier functions.
414 ::
416 -- / L \ 2 2
417 radius**2 = - > | --- | ln |Z|
418 --d \ 2pi /
419 """
420 r2 = -np.dot(self.largeunitcell_cc.diagonal()**2 / (2 * pi)**2,
421 np.log(abs(self.Z_dww[:3].diagonal(0, 1, 2))**2))
422 return np.sqrt(r2)
424 def get_spectral_weight(self, w):
425 return abs(self.V_knw[:, :, w])**2 / self.Nk
427 def get_pdos(self, w, energies, width):
428 """Projected density of states (PDOS).
430 Returns the (PDOS) for Wannier function ``w``. The calculation
431 is performed over the energy grid specified in energies. The
432 PDOS is produced as a sum of Gaussians centered at the points
433 of the energy grid and with the specified width.
434 """
435 spec_kn = self.get_spectral_weight(w)
436 dos = np.zeros(len(energies))
437 for k, spec_n in enumerate(spec_kn):
438 eig_n = self.calc.get_eigenvalues(kpt=k, spin=self.spin)
439 for weight, eig in zip(spec_n, eig_n):
440 # Add gaussian centered at the eigenvalue
441 x = ((energies - eig) / width)**2
442 dos += weight * np.exp(-x.clip(0., 40.)) / (sqrt(pi) * width)
443 return dos
445 def translate(self, w, R):
446 """Translate the w'th Wannier function
448 The distance vector R = [n1, n2, n3], is in units of the basis
449 vectors of the small cell.
450 """
451 for kpt_c, U_ww in zip(self.kpt_kc, self.U_kww):
452 U_ww[:, w] *= np.exp(2.j * pi * np.dot(np.array(R), kpt_c))
453 self.update()
455 def translate_to_cell(self, w, cell):
456 """Translate the w'th Wannier function to specified cell"""
457 scaled_c = np.angle(self.Z_dww[:3, w, w]) * self.kptgrid / (2 * pi)
458 trans = np.array(cell) - np.floor(scaled_c)
459 self.translate(w, trans)
461 def translate_all_to_cell(self, cell=[0, 0, 0]):
462 r"""Translate all Wannier functions to specified cell.
464 Move all Wannier orbitals to a specific unit cell. There
465 exists an arbitrariness in the positions of the Wannier
466 orbitals relative to the unit cell. This method can move all
467 orbitals to the unit cell specified by ``cell``. For a
468 `\Gamma`-point calculation, this has no effect. For a
469 **k**-point calculation the periodicity of the orbitals are
470 given by the large unit cell defined by repeating the original
471 unitcell by the number of **k**-points in each direction. In
472 this case it is useful to move the orbitals away from the
473 boundaries of the large cell before plotting them. For a bulk
474 calculation with, say 10x10x10 **k** points, one could move
475 the orbitals to the cell [2,2,2]. In this way the pbc
476 boundary conditions will not be noticed.
477 """
478 scaled_wc = (np.angle(self.Z_dww[:3].diagonal(0, 1, 2)).T *
479 self.kptgrid / (2 * pi))
480 trans_wc = np.array(cell)[None] - np.floor(scaled_wc)
481 for kpt_c, U_ww in zip(self.kpt_kc, self.U_kww):
482 U_ww *= np.exp(2.j * pi * np.dot(trans_wc, kpt_c))
483 self.update()
485 def distances(self, R):
486 """Relative distances between centers.
488 Returns a matrix with the distances between different Wannier centers.
489 R = [n1, n2, n3] is in units of the basis vectors of the small cell
490 and allows one to measure the distance with centers moved to a
491 different small cell.
492 The dimension of the matrix is [Nw, Nw].
493 """
494 Nw = self.nwannier
495 cen = self.get_centers()
496 r1 = cen.repeat(Nw, axis=0).reshape(Nw, Nw, 3)
497 r2 = cen.copy()
498 for i in range(3):
499 r2 += self.unitcell_cc[i] * R[i]
501 r2 = np.swapaxes(r2.repeat(Nw, axis=0).reshape(Nw, Nw, 3), 0, 1)
502 return np.sqrt(np.sum((r1 - r2)**2, axis=-1))
504 def get_hopping(self, R):
505 """Returns the matrix H(R)_nm=<0,n|H|R,m>.
507 ::
509 1 _ -ik.R
510 H(R) = <0,n|H|R,m> = --- >_ e H(k)
511 Nk k
513 where R is the cell-distance (in units of the basis vectors of
514 the small cell) and n,m are indices of the Wannier functions.
515 """
516 H_ww = np.zeros([self.nwannier, self.nwannier], complex)
517 for k, kpt_c in enumerate(self.kpt_kc):
518 phase = np.exp(-2.j * pi * np.dot(np.array(R), kpt_c))
519 H_ww += self.get_hamiltonian(k) * phase
520 return H_ww / self.Nk
522 def get_hamiltonian(self, k=0):
523 """Get Hamiltonian at existing k-vector of index k
525 ::
527 dag
528 H(k) = V diag(eps ) V
529 k k k
530 """
531 eps_n = self.calc.get_eigenvalues(kpt=k, spin=self.spin)[:self.nbands]
532 return np.dot(dag(self.V_knw[k]) * eps_n, self.V_knw[k])
534 def get_hamiltonian_kpoint(self, kpt_c):
535 """Get Hamiltonian at some new arbitrary k-vector
537 ::
539 _ ik.R
540 H(k) = >_ e H(R)
541 R
543 Warning: This method moves all Wannier functions to cell (0, 0, 0)
544 """
545 if self.verbose:
546 print('Translating all Wannier functions to cell (0, 0, 0)')
547 self.translate_all_to_cell()
548 max = (self.kptgrid - 1) // 2
549 N1, N2, N3 = max
550 Hk = np.zeros([self.nwannier, self.nwannier], complex)
551 for n1 in range(-N1, N1 + 1):
552 for n2 in range(-N2, N2 + 1):
553 for n3 in range(-N3, N3 + 1):
554 R = np.array([n1, n2, n3], float)
555 hop_ww = self.get_hopping(R)
556 phase = np.exp(+2.j * pi * np.dot(R, kpt_c))
557 Hk += hop_ww * phase
558 return Hk
560 def get_function(self, index, repeat=None):
561 r"""Get Wannier function on grid.
563 Returns an array with the funcion values of the indicated Wannier
564 function on a grid with the size of the *repeated* unit cell.
566 For a calculation using **k**-points the relevant unit cell for
567 eg. visualization of the Wannier orbitals is not the original unit
568 cell, but rather a larger unit cell defined by repeating the
569 original unit cell by the number of **k**-points in each direction.
570 Note that for a `\Gamma`-point calculation the large unit cell
571 coinsides with the original unit cell.
572 The large unitcell also defines the periodicity of the Wannier
573 orbitals.
575 ``index`` can be either a single WF or a coordinate vector in terms
576 of the WFs.
577 """
579 # Default size of plotting cell is the one corresponding to k-points.
580 if repeat is None:
581 repeat = self.kptgrid
582 N1, N2, N3 = repeat
584 dim = self.calc.get_number_of_grid_points()
585 largedim = dim * [N1, N2, N3]
587 wanniergrid = np.zeros(largedim, dtype=complex)
588 for k, kpt_c in enumerate(self.kpt_kc):
589 # The coordinate vector of wannier functions
590 if isinstance(index, int):
591 vec_n = self.V_knw[k, :, index]
592 else:
593 vec_n = np.dot(self.V_knw[k], index)
595 wan_G = np.zeros(dim, complex)
596 for n, coeff in enumerate(vec_n):
597 wan_G += coeff * self.calc.get_pseudo_wave_function(
598 n, k, self.spin, pad=True)
600 # Distribute the small wavefunction over large cell:
601 for n1 in range(N1):
602 for n2 in range(N2):
603 for n3 in range(N3): # sign?
604 e = np.exp(-2.j * pi * np.dot([n1, n2, n3], kpt_c))
605 wanniergrid[n1 * dim[0]:(n1 + 1) * dim[0],
606 n2 * dim[1]:(n2 + 1) * dim[1],
607 n3 * dim[2]:(n3 + 1) * dim[2]] += e * wan_G
609 # Normalization
610 wanniergrid /= np.sqrt(self.Nk)
611 return wanniergrid
613 def write_cube(self, index, fname, repeat=None, real=True):
614 """Dump specified Wannier function to a cube file"""
615 from ase.io import write
617 # Default size of plotting cell is the one corresponding to k-points.
618 if repeat is None:
619 repeat = self.kptgrid
620 atoms = self.calc.get_atoms() * repeat
621 func = self.get_function(index, repeat)
623 # Handle separation of complex wave into real parts
624 if real:
625 if self.Nk == 1:
626 func *= np.exp(-1.j * np.angle(func.max()))
627 if 0:
628 assert max(abs(func.imag).flat) < 1e-4
629 func = func.real
630 else:
631 func = abs(func)
632 else:
633 phase_fname = fname.split('.')
634 phase_fname.insert(1, 'phase')
635 phase_fname = '.'.join(phase_fname)
636 write(phase_fname, atoms, data=np.angle(func), format='cube')
637 func = abs(func)
639 write(fname, atoms, data=func, format='cube')
641 def localize(self, step=0.25, tolerance=1e-08,
642 updaterot=True, updatecoeff=True):
643 """Optimize rotation to give maximal localization"""
644 md_min(self, step, tolerance, verbose=self.verbose,
645 updaterot=updaterot, updatecoeff=updatecoeff)
647 def get_functional_value(self):
648 """Calculate the value of the spread functional.
650 ::
652 Tr[|ZI|^2]=sum(I)sum(n) w_i|Z_(i)_nn|^2,
654 where w_i are weights."""
655 a_d = np.sum(np.abs(self.Z_dww.diagonal(0, 1, 2))**2, axis=1)
656 return np.dot(a_d, self.weight_d).real
658 def get_gradients(self):
659 # Determine gradient of the spread functional.
660 #
661 # The gradient for a rotation A_kij is::
662 #
663 # dU = dRho/dA_{k,i,j} = sum(I) sum(k')
664 # + Z_jj Z_kk',ij^* - Z_ii Z_k'k,ij^*
665 # - Z_ii^* Z_kk',ji + Z_jj^* Z_k'k,ji
666 #
667 # The gradient for a change of coefficients is::
668 #
669 # dRho/da^*_{k,i,j} = sum(I) [[(Z_0)_{k} V_{k'} diag(Z^*) +
670 # (Z_0_{k''})^d V_{k''} diag(Z)] *
671 # U_k^d]_{N+i,N+j}
672 #
673 # where diag(Z) is a square,diagonal matrix with Z_nn in the diagonal,
674 # k' = k + dk and k = k'' + dk.
675 #
676 # The extra degrees of freedom chould be kept orthonormal to the fixed
677 # space, thus we introduce lagrange multipliers, and minimize instead::
678 #
679 # Rho_L=Rho- sum_{k,n,m} lambda_{k,nm} <c_{kn}|c_{km}>
680 #
681 # for this reason the coefficient gradients should be multiplied
682 # by (1 - c c^d).
684 Nb = self.nbands
685 Nw = self.nwannier
687 dU = []
688 dC = []
689 for k in range(self.Nk):
690 M = self.fixedstates_k[k]
691 L = self.edf_k[k]
692 U_ww = self.U_kww[k]
693 C_ul = self.C_kul[k]
694 Utemp_ww = np.zeros((Nw, Nw), complex)
695 Ctemp_nw = np.zeros((Nb, Nw), complex)
697 for d, weight in enumerate(self.weight_d):
698 if abs(weight) < 1.0e-6:
699 continue
701 Z_knn = self.Z_dknn[d]
702 diagZ_w = self.Z_dww[d].diagonal()
703 Zii_ww = np.repeat(diagZ_w, Nw).reshape(Nw, Nw)
704 k1 = self.kklst_dk[d, k]
705 k2 = self.invkklst_dk[d, k]
706 V_knw = self.V_knw
707 Z_kww = self.Z_dkww[d]
709 if L > 0:
710 Ctemp_nw += weight * np.dot(
711 np.dot(Z_knn[k], V_knw[k1]) * diagZ_w.conj() +
712 np.dot(dag(Z_knn[k2]), V_knw[k2]) * diagZ_w,
713 dag(U_ww))
715 temp = Zii_ww.T * Z_kww[k].conj() - Zii_ww * Z_kww[k2].conj()
716 Utemp_ww += weight * (temp - dag(temp))
717 dU.append(Utemp_ww.ravel())
718 if L > 0:
719 # Ctemp now has same dimension as V, the gradient is in the
720 # lower-right (Nb-M) x L block
721 Ctemp_ul = Ctemp_nw[M:, M:]
722 G_ul = Ctemp_ul - np.dot(np.dot(C_ul, dag(C_ul)), Ctemp_ul)
723 dC.append(G_ul.ravel())
725 return np.concatenate(dU + dC)
727 def step(self, dX, updaterot=True, updatecoeff=True):
728 # dX is (A, dC) where U->Uexp(-A) and C->C+dC
729 Nw = self.nwannier
730 Nk = self.Nk
731 M_k = self.fixedstates_k
732 L_k = self.edf_k
733 if updaterot:
734 A_kww = dX[:Nk * Nw**2].reshape(Nk, Nw, Nw)
735 for U, A in zip(self.U_kww, A_kww):
736 H = -1.j * A.conj()
737 epsilon, Z = np.linalg.eigh(H)
738 # Z contains the eigenvectors as COLUMNS.
739 # Since H = iA, dU = exp(-A) = exp(iH) = ZDZ^d
740 dU = np.dot(Z * np.exp(1.j * epsilon), dag(Z))
741 if U.dtype == float:
742 U[:] = np.dot(U, dU).real
743 else:
744 U[:] = np.dot(U, dU)
746 if updatecoeff:
747 start = 0
748 for C, unocc, L in zip(self.C_kul, self.nbands - M_k, L_k):
749 if L == 0 or unocc == 0:
750 continue
751 Ncoeff = L * unocc
752 deltaC = dX[Nk * Nw**2 + start: Nk * Nw**2 + start + Ncoeff]
753 C += deltaC.reshape(unocc, L)
754 gram_schmidt(C)
755 start += Ncoeff
757 self.update()