Coverage for gpaw/xc/noncollinear.py: 98%

1import numpy as np 

2from scipy.linalg import eigh 

3 

4from gpaw.lcao.eigensolver import DirectLCAO 

5 

6 

7class NonCollinearLDAKernel: 

8 name = 'LDA' 

9 type = 'LDA' 

10 

11 def __init__(self, kernel): 

12 self.kernel = kernel 

13 

14 def calculate(self, e_g, n_sg, v_sg): 

15 n_g = n_sg[0] 

16 m_vg = n_sg[1:4] 

17 m_g = (m_vg**2).sum(0)**0.5 

18 nnew_sg = np.empty((2,) + n_g.shape) 

19 nnew_sg[:] = n_g 

20 nnew_sg[0] += m_g 

21 nnew_sg[1] -= m_g 

22 nnew_sg *= 0.5 

23 vnew_sg = np.zeros_like(nnew_sg) 

24 self.kernel.calculate(e_g, nnew_sg, vnew_sg) 

25 v_sg[0] += 0.5 * vnew_sg.sum(0) 

26 vnew_sg /= np.where(m_g < 1e-15, 1, m_g) 

27 v_sg[1:4] += 0.5 * vnew_sg[0] * m_vg 

28 v_sg[1:4] -= 0.5 * vnew_sg[1] * m_vg 

29 

30 

31class NonCollinearLCAOEigensolver(DirectLCAO): 

32 def iterate(self, ham, wfs): 

33 wfs.timer.start('LCAO eigensolver') 

34 

35 wfs.timer.start('Potential matrix') 

36 Vt_xdMM = [wfs.basis_functions.calculate_potential_matrices(vt_G) 

37 for vt_G in ham.vt_xG] 

38 wfs.timer.stop('Potential matrix') 

39 

40 for kpt in wfs.kpt_u: 

41 self.iterate_one_k_point(ham, wfs, kpt, Vt_xdMM) 

42 

43 wfs.timer.stop('LCAO eigensolver') 

44 

45 def iterate_one_k_point(self, ham, wfs, kpt, Vt_xdMM): 

46 assert wfs.gd.comm.size == 1, 'No quite sure this works!' 

47 if wfs.bd.comm.size > 1 and wfs.bd.strided: 

48 raise NotImplementedError 

49 

50 H_xMM = [] 

51 for x in range(4): 

52 kpt.s = x 

53 H_MM = self.calculate_hamiltonian_matrix(ham, wfs, kpt, Vt_xdMM[x], 

54 root=0, 

55 add_kinetic=(x == 0)) 

56 H_xMM.append(H_MM) 

57 kpt.s = None 

58 

59 S_MM = wfs.S_qMM[kpt.q] 

60 M = len(S_MM) 

61 S2_MM = np.zeros((2 * M, 2 * M), complex) 

62 H2_MM = np.zeros((2 * M, 2 * M), complex) 

63 

64 S2_MM[:M, :M] = S_MM 

65 S2_MM[M:, M:] = S_MM 

66 

67 # We are ignoring x and y here. 

68 # See gpaw.new.lcao.hamiltonian.NonCollinearHamiltonianMatrixCalculator 

69 # for the correct way! 

70 H2_MM[:M, :M] = H_xMM[0] + H_xMM[3] 

71 H2_MM[M:, M:] = H_xMM[0] - H_xMM[3] 

72 

73 kpt.eps_n = np.empty(2 * wfs.bd.mynbands) 

74 

75 diagonalization_string = repr(self.diagonalizer) 

76 wfs.timer.start(diagonalization_string) 

77 kpt.eps_n, V_MM = eigh(H2_MM, S2_MM) 

78 kpt.C_nM = V_MM.T.conj() 

79 wfs.timer.stop(diagonalization_string) 

80 kpt.C_nM.shape = (wfs.bd.mynbands * 4, M)