Coverage for gpaw/directmin/functional/fdpw/ks.py: 85%

67 statements  

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

1""" 

2Potentials for orbital density dependent energy functionals 

3""" 

4 

5import numpy as np 

6 

7from gpaw.directmin.tools import d_matrix 

8from gpaw.utilities import unpack_hermitian 

9 

10 

11class KSFDPW: 

12 """ 

13 KS-DFT 

14 """ 

15 def __init__(self, wfs, dens, ham): 

16 

17 self.name = 'Zero' 

18 self.n_kps = wfs.kd.nibzkpts 

19 self.grad = {} 

20 self.total_sic = 0.0 

21 self.eks = 0.0 

22 self.changedocc = 0 

23 self.dens = dens 

24 self.ham = ham 

25 

26 def get_energy_and_gradients_kpt( 

27 self, wfs, kpt, grad_knG=None, U_k=None, add_grad=False, ham=None): 

28 

29 k = self.n_kps * kpt.s + kpt.q 

30 nbands = wfs.bd.nbands 

31 if U_k is not None: 

32 assert U_k[k].shape[0] == nbands 

33 

34 wfs.timer.start('e/g grid calculations') 

35 self.grad[k] = wfs.empty(nbands, q=kpt.q) 

36 wfs.apply_pseudo_hamiltonian(kpt, ham, kpt.psit_nG, self.grad[k]) 

37 

38 c_axi = {} 

39 for a, P_xi in kpt.P_ani.items(): 

40 try: 

41 dH_ii = unpack_hermitian(ham.dH_asp[a][kpt.s]) 

42 except AttributeError: 

43 dH_ii = ham.potential.dH_asii[a][kpt.s] 

44 c_xi = np.dot(P_xi, dH_ii) 

45 c_axi[a] = c_xi 

46 

47 # not sure about this: 

48 ham.xc.add_correction(kpt, kpt.psit_nG, self.grad[k], 

49 kpt.P_ani, c_axi, n_x=None, 

50 calculate_change=False) 

51 # add projectors to the H|psi_i> 

52 wfs.pt.add(self.grad[k], c_axi, kpt.q) 

53 # scale with occupation numbers 

54 for i, f in enumerate(kpt.f_n): 

55 self.grad[k][i] *= f 

56 

57 if add_grad: 

58 if U_k is not None: 

59 grad_knG[k] += \ 

60 np.tensordot(U_k[k].conj(), self.grad[k], axes=1) 

61 else: 

62 grad_knG[k] += self.grad[k] 

63 else: 

64 if U_k is not None: 

65 self.grad[k][:] = np.tensordot(U_k[k].conj(), 

66 self.grad[k], axes=1) 

67 

68 wfs.timer.stop('e/g grid calculations') 

69 

70 return 0.0 

71 

72 def get_energy_and_gradients_inner_loop( 

73 self, wfs, kpt, a_mat, evals, evec, ham=None, exstate=True): 

74 

75 if not exstate: 

76 raise RuntimeError('Attempting to optimize unitary-invariant ' 

77 'energy in occupied-occupied rotation space.') 

78 

79 ndim = wfs.bd.nbands 

80 

81 k = self.n_kps * kpt.s + kpt.q 

82 self.grad[k] = np.zeros_like(kpt.psit_nG[:ndim]) 

83 e_sic = self.get_energy_and_gradients_kpt(wfs, kpt, ham=ham) 

84 wfs.timer.start('Unitary gradients') 

85 l_odd = wfs.integrate(kpt.psit_nG[:ndim], self.grad[k][:ndim], True) 

86 f = np.ones(ndim) 

87 indz = np.absolute(l_odd) > 1.0e-4 

88 l_c = 2.0 * l_odd[indz] 

89 l_odd = f[:, np.newaxis] * l_odd.T.conj() - f * l_odd 

90 kappa = np.max(np.absolute(l_odd[indz]) / np.absolute(l_c)) 

91 

92 if a_mat is None: 

93 wfs.timer.stop('Unitary gradients') 

94 return 2.0 * l_odd.T.conj(), e_sic, kappa 

95 else: 

96 g_mat = evec.T.conj() @ l_odd.T.conj() @ evec 

97 g_mat = g_mat * d_matrix(evals) 

98 g_mat = evec @ g_mat @ evec.T.conj() 

99 for i in range(g_mat.shape[0]): 

100 g_mat[i][i] *= 0.5 

101 wfs.timer.stop('Unitary gradients') 

102 if a_mat.dtype == float: 

103 g_mat = g_mat.real 

104 return 2.0 * g_mat, e_sic, kappa