Coverage for /builds/kinetik161/ase/ase/dft/wannierstate.py: 63.79%

1import numpy as np 

2from scipy.linalg import qr 

3 

4 

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) 

14 

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))) 

19 

20 return ortho_m 

21 

22 

23def _empty(): 

24 return np.empty(0, complex) 

25 

26 

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 

33 

34 def _zeros(self): 

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

36 

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) 

47 

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) 

60 

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) 

65 

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) 

71 

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

73 from ase.dft.wannier import scdm 

74 

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

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

77 kpt=0, spin=0) 

78 Ng = ps.size 

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

80 dtype=np.complex128) 

81 for k in range(self.Nk): 

82 for n in range(self.Nb): 

83 pseudo_nkG[n, k] = \ 

84 calc.get_pseudo_wave_function( 

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

86 

87 # Use initial guess to determine U and C 

88 C_kul, U_kww = scdm(pseudo_nkG, 

89 kpts=kpt_kc, 

90 fixed_k=self.fixedstates_k, 

91 Nw=self.Nw) 

92 return WannierState(C_kul, U_kww) 

93 

94 

95class WannierState: 

96 def __init__(self, C_kul, U_kww): 

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

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

99 self.U_kww = U_kww