Coverage for gpaw/lcao/pwf2.py: 92%

84 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-07-08 00:17 +0000

1import numpy as np 

2from ase.units import Hartree 

3 

4from gpaw.lcao.projected_wannier import condition_number, dots, get_bfs 

5from gpaw.utilities import unpack_hermitian 

6from gpaw.utilities.tools import tri2full 

7 

8 

9def get_lcao_xc(calc, P_aqMi, bfs=None, spin=0): 

10 nq = len(calc.wfs.kd.ibzk_qc) 

11 nao = calc.wfs.setups.nao 

12 dtype = calc.wfs.dtype 

13 if bfs is None: 

14 bfs = get_bfs(calc) 

15 

16 if calc.density.nt_sg is None: 

17 calc.density.interpolate_pseudo_density() 

18 nt_sg = calc.density.nt_sg 

19 vxct_sg = calc.density.finegd.zeros(calc.wfs.nspins) 

20 calc.hamiltonian.xc.calculate(calc.density.finegd, nt_sg, vxct_sg) 

21 vxct_G = calc.wfs.gd.zeros() 

22 calc.hamiltonian.restrict_and_collect(vxct_sg[spin], vxct_G) 

23 Vxc_qMM = np.zeros((nq, nao, nao), dtype) 

24 for q, Vxc_MM in enumerate(Vxc_qMM): 

25 bfs.calculate_potential_matrix(vxct_G, Vxc_MM, q) 

26 tri2full(Vxc_MM, 'L') 

27 

28 # Add atomic PAW corrections 

29 for a, P_qMi in P_aqMi.items(): 

30 D_sp = calc.density.D_asp[a][:] 

31 H_sp = np.zeros_like(D_sp) 

32 calc.hamiltonian.xc.calculate_paw_correction(calc.wfs.setups[a], 

33 D_sp, H_sp) 

34 H_ii = unpack_hermitian(H_sp[spin]) 

35 for Vxc_MM, P_Mi in zip(Vxc_qMM, P_qMi): 

36 Vxc_MM += dots(P_Mi, H_ii, P_Mi.T.conj()) 

37 return Vxc_qMM * Hartree 

38 

39 

40class LCAOwrap: 

41 def __init__(self, calc, spin=0): 

42 assert calc.wfs.gd.comm.size == 1 

43 assert calc.wfs.kd.comm.size == 1 

44 assert calc.wfs.bd.comm.size == 1 

45 

46 from gpaw.lcao.tools import get_lcao_hamiltonian 

47 H_skMM, S_kMM = get_lcao_hamiltonian(calc) 

48 

49 self.calc = calc 

50 self.dtype = calc.wfs.dtype 

51 self.spin = spin 

52 self.H_qww = H_skMM[spin] 

53 self.S_qww = S_kMM 

54 self.P_aqwi = calc.wfs.P_aqMi 

55 self.Nw = self.S_qww.shape[-1] 

56 

57 for S in self.S_qww: 

58 print('Condition number: %0.1e' % condition_number(S)) 

59 

60 def get_hamiltonian(self, q=0, indices=None): 

61 if indices is None: 

62 return self.H_qww[q] 

63 else: 

64 return self.H_qww[q].take(indices, 0).take(indices, 1) 

65 

66 def get_overlap(self, q=0, indices=None): 

67 if indices is None: 

68 return self.S_qww[q] 

69 else: 

70 return self.S_qww[q].take(indices, 0).take(indices, 1) 

71 

72 def get_projections(self, q=0, indices=None): 

73 if indices is None: 

74 return {a: P_qwi[q] for a, P_qwi in self.P_aqwi.items()} 

75 else: 

76 return {a: P_qwi[q].take(indices, 0) 

77 for a, P_qwi in self.P_aqwi.items()} 

78 

79 def get_orbitals(self, q=-1, indices=None): 

80 assert q == -1 

81 if indices is None: 

82 indices = range(self.Nw) 

83 Ni = len(indices) 

84 C_wM = np.zeros((Ni, self.Nw), self.dtype) 

85 for i, C_M in zip(indices, C_wM): 

86 C_M[i] = 1.0 

87 w_wG = self.calc.wfs.gd.zeros(Ni, dtype=self.dtype) 

88 self.calc.wfs.basis_functions.lcao_to_grid(C_wM, w_wG, q=-1) 

89 return w_wG 

90 

91 def get_Fcore(self, q=0, indices=None): 

92 if indices is None: 

93 Fcore_ww = np.zeros_like(self.H_qww[q]) 

94 else: 

95 Fcore_ww = np.zeros((len(indices), len(indices))) 

96 for a, P_wi in self.get_projections(q, indices).items(): 

97 if self.calc.wfs.setups[a].type != 'ghost': 

98 X_ii = unpack_hermitian(self.calc.wfs.setups[a].X_p) 

99 Fcore_ww -= dots(P_wi, X_ii, P_wi.T.conj()) 

100 return Fcore_ww * Hartree 

101 

102 def get_xc(self, q=0, indices=None): 

103 if not hasattr(self, 'Vxc_qww'): 

104 self.Vxc_qww = get_lcao_xc(self.calc, self.P_aqwi, 

105 bfs=self.calc.wfs.basis_functions, 

106 spin=self.spin) 

107 if indices is None: 

108 return self.Vxc_qww[q] 

109 else: 

110 return self.Vxc_qww[q].take(indices, 0).take(indices, 1)