r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1import numpy as np

2from scipy.linalg import qr

5def random_orthogonal_matrix(dim, rng=np.random, real=False):

6 """Generate uniformly distributed random orthogonal matrices"""

7 if real:

8 from scipy.stats import special_ortho_group

9 ortho_m = special_ortho_group.rvs(dim=dim, random_state=rng)

10 else:

11 # The best method but not supported on old systems

12 # from scipy.stats import unitary_group

13 # ortho_m = unitary_group.rvs(dim=dim, random_state=rng)

15 # Alternative method from https://stackoverflow.com/questions/38426349

16 H = rng.random((dim, dim))

17 Q, R = qr(H)

18 ortho_m = Q @ np.diag(np.sign(np.diag(R)))

20 return ortho_m

23def _empty():

24 return np.empty(0, complex)

27class WannierSpec:

28 def __init__(self, Nk, Nw, Nb, fixedstates_k):

29 self.Nk = Nk

30 self.Nw = Nw

31 self.Nb = Nb

32 self.fixedstates_k = fixedstates_k

34 def _zeros(self):

35 return np.zeros((self.Nk, self.Nw, self.Nw), complex)

37 def bloch(self, edf_k):

38 U_kww = self._zeros()

39 C_kul = []

40 for U, M, L in zip(U_kww, self.fixedstates_k, edf_k):

41 U[:] = np.identity(self.Nw, complex)

42 if L > 0:

43 C_kul.append(np.identity(self.Nb - M, complex)[:, :L])

44 else:

45 C_kul.append(_empty())

46 return WannierState(C_kul, U_kww)

48 def random(self, rng, edf_k):

49 # Set U and C to random (orthogonal) matrices

50 U_kww = self._zeros()

51 C_kul = []

52 for U, M, L in zip(U_kww, self.fixedstates_k, edf_k):

53 U[:] = random_orthogonal_matrix(self.Nw, rng, real=False)

54 if L > 0:

55 C_kul.append(random_orthogonal_matrix(

56 self.Nb - M, rng=rng, real=False)[:, :L])

57 else:

58 C_kul.append(_empty())

59 return WannierState(C_kul, U_kww)

61 def initial_orbitals(self, calc, orbitals, kptgrid, edf_k, spin):

62 C_kul, U_kww = calc.initial_wannier(

63 orbitals, kptgrid, self.fixedstates_k, edf_k, spin, self.Nb)

64 return WannierState(C_kul, U_kww)

66 def initial_wannier(self, calc, method, kptgrid, edf_k, spin):

67 C_kul, U_kww = calc.initial_wannier(

68 method, kptgrid, self.fixedstates_k,

69 edf_k, spin, self.Nb)

70 return WannierState(C_kul, U_kww)

72 def scdm(self, calc, kpt_kc, spin):

73 from ase.dft.wannier import scdm

74 # get the size of the grid and check if there are Nw bands:

75 ps = calc.get_pseudo_wave_function(band=self.Nw,

76 kpt=0, spin=0)

77 Ng = ps.size

78 pseudo_nkG = np.zeros((self.Nb, self.Nk, Ng),

79 dtype=np.complex128)

80 for k in range(self.Nk):

81 for n in range(self.Nb):

82 pseudo_nkG[n, k] = \

83 calc.get_pseudo_wave_function(

84 band=n, kpt=k, spin=spin).ravel()

86 # Use initial guess to determine U and C

87 C_kul, U_kww = scdm(pseudo_nkG,

88 kpts=kpt_kc,

89 fixed_k=self.fixedstates_k,

90 Nw=self.Nw)

91 return WannierState(C_kul, U_kww)

94class WannierState:

95 def __init__(self, C_kul, U_kww):

96 # Number of u is not always the same, so C_kul is ragged

97 self.C_kul = [C_ul.astype(complex) for C_ul in C_kul]

98 self.U_kww = U_kww