1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
|
import numpy as np
from scipy.linalg import qr
def random_orthogonal_matrix(dim, rng=np.random, real=False):
"""Generate uniformly distributed random orthogonal matrices"""
if real:
from scipy.stats import special_ortho_group
ortho_m = special_ortho_group.rvs(dim=dim, random_state=rng)
else:
# The best method but not supported on old systems
# from scipy.stats import unitary_group
# ortho_m = unitary_group.rvs(dim=dim, random_state=rng)
# Alternative method from https://stackoverflow.com/questions/38426349
H = rng.random((dim, dim))
Q, R = qr(H)
ortho_m = Q @ np.diag(np.sign(np.diag(R)))
return ortho_m
def _empty():
return np.empty(0, complex)
class WannierSpec:
def __init__(self, Nk, Nw, Nb, fixedstates_k):
self.Nk = Nk
self.Nw = Nw
self.Nb = Nb
self.fixedstates_k = fixedstates_k
def _zeros(self):
return np.zeros((self.Nk, self.Nw, self.Nw), complex)
def bloch(self, edf_k):
U_kww = self._zeros()
C_kul = []
for U, M, L in zip(U_kww, self.fixedstates_k, edf_k):
U[:] = np.identity(self.Nw, complex)
if L > 0:
C_kul.append(np.identity(self.Nb - M, complex)[:, :L])
else:
C_kul.append(_empty())
return WannierState(C_kul, U_kww)
def random(self, rng, edf_k):
# Set U and C to random (orthogonal) matrices
U_kww = self._zeros()
C_kul = []
for U, M, L in zip(U_kww, self.fixedstates_k, edf_k):
U[:] = random_orthogonal_matrix(self.Nw, rng, real=False)
if L > 0:
C_kul.append(random_orthogonal_matrix(
self.Nb - M, rng=rng, real=False)[:, :L])
else:
C_kul.append(_empty())
return WannierState(C_kul, U_kww)
def initial_orbitals(self, calc, orbitals, kptgrid, edf_k, spin):
C_kul, U_kww = calc.initial_wannier(
orbitals, kptgrid, self.fixedstates_k, edf_k, spin, self.Nb)
return WannierState(C_kul, U_kww)
def initial_wannier(self, calc, method, kptgrid, edf_k, spin):
C_kul, U_kww = calc.initial_wannier(
method, kptgrid, self.fixedstates_k,
edf_k, spin, self.Nb)
return WannierState(C_kul, U_kww)
def scdm(self, calc, kpt_kc, spin):
from ase.dft.wannier import scdm
# get the size of the grid and check if there are Nw bands:
ps = calc.get_pseudo_wave_function(band=self.Nw,
kpt=0, spin=0)
Ng = ps.size
pseudo_nkG = np.zeros((self.Nb, self.Nk, Ng),
dtype=np.complex128)
for k in range(self.Nk):
for n in range(self.Nb):
pseudo_nkG[n, k] = \
calc.get_pseudo_wave_function(
band=n, kpt=k, spin=spin).ravel()
# Use initial guess to determine U and C
C_kul, U_kww = scdm(pseudo_nkG,
kpts=kpt_kc,
fixed_k=self.fixedstates_k,
Nw=self.Nw)
return WannierState(C_kul, U_kww)
class WannierState:
def __init__(self, C_kul, U_kww):
# Number of u is not always the same, so C_kul is ragged
self.C_kul = [C_ul.astype(complex) for C_ul in C_kul]
self.U_kww = U_kww
|
