Coverage for gpaw/forces.py: 99%

1import numpy as np 

2 

3from ase.units import Hartree, Bohr 

4 

5from gpaw.directmin.tools import get_n_occ 

6from gpaw.utilities import unpack_hermitian 

7from gpaw.xc.hybrid import HybridXCBase 

8 

9 

10def calculate_forces(wfs, dens, ham, log=None): 

11 """Return the atomic forces.""" 

12 

13 assert not isinstance(ham.xc, HybridXCBase) 

14 assert not ham.xc.name.startswith('GLLB') 

15 

16 func_name = None 

17 if hasattr(wfs.eigensolver, 'dm_helper'): 

18 func_name = getattr(wfs.eigensolver.dm_helper.func, 'name', None) 

19 elif hasattr(wfs.eigensolver, 'odd'): 

20 func_name = getattr(wfs.eigensolver.odd, 'name', None) 

21 if func_name == 'PZ-SIC': 

22 if wfs.mode == 'fd' or wfs.mode == 'pw': 

23 return calculate_forces_using_non_diag_lagr_matrix( 

24 wfs, dens, ham, log) 

25 elif wfs.mode == 'lcao': 

26 for kpt in wfs.kpt_u: 

27 kpt.rho_MM = \ 

28 wfs.calculate_density_matrix(kpt.f_n, kpt.C_nM) 

29 

30 natoms = len(wfs.setups) 

31 

32 # Force from projector functions (and basis set): 

33 F_wfs_av = np.zeros((natoms, 3)) 

34 wfs.calculate_forces(ham, F_wfs_av) 

35 wfs.gd.comm.sum(F_wfs_av, 0) 

36 F_ham_av = np.zeros((natoms, 3)) 

37 

38 try: 

39 # ODD functionals need force corrections for each spin 

40 correction = ham.xc.setup_force_corrections 

41 except AttributeError: 

42 pass 

43 else: 

44 correction(F_ham_av) 

45 

46 ham.calculate_forces(dens, F_ham_av) 

47 

48 F_av = F_ham_av + F_wfs_av 

49 wfs.world.broadcast(F_av, 0) 

50 

51 if func_name == 'PZ-SIC' and wfs.mode == 'lcao': 

52 F_av += wfs.eigensolver.dm_helper.func.get_odd_corrections_to_forces( 

53 wfs, dens) 

54 for kpt in wfs.kpt_u: 

55 # need to re-set rho_MM otherwise it will be used 

56 # it's probably better to in wfs.reset, but 

57 # when position changes wfs.reset is not called 

58 kpt.rho_MM = None 

59 

60 F_av = wfs.kd.symmetry.symmetrize_forces(F_av) 

61 

62 if log: 

63 log('\nForces in eV/Ang:') 

64 c = Hartree / Bohr 

65 for a, setup in enumerate(wfs.setups): 

66 log('%3d %-2s %10.5f %10.5f %10.5f' % 

67 ((a, setup.symbol) + tuple(F_av[a] * c))) 

68 log() 

69 

70 return F_av 

71 

72 

73def calculate_forces_using_non_diag_lagr_matrix(wfs, dens, ham, log=None): 

74 

75 natoms = len(wfs.setups) 

76 F_wfs_av = np.zeros((natoms, 3)) 

77 esolv = wfs.eigensolver 

78 

79 grad_knG = esolv.get_gradients_2(ham, wfs) 

80 if 'SIC' in esolv.odd.name: 

81 for kpt in wfs.kpt_u: 

82 esolv.odd.get_energy_and_gradients_kpt( 

83 wfs, kpt, grad_knG, esolv.iloop.U_k, add_grad=True) 

84 for kpt in wfs.kpt_u: 

85 k = esolv.n_kps * kpt.s + kpt.q 

86 n_occ = get_n_occ(kpt)[0] 

87 

88 lamb = wfs.integrate( 

89 kpt.psit_nG[:n_occ], grad_knG[k][:n_occ], True) 

90 

91 P_ani = kpt.P_ani 

92 dP_aniv = wfs.pt.dict(n_occ, derivative=True) 

93 wfs.pt.derivative(kpt.psit_nG[:n_occ], dP_aniv, kpt.q) 

94 dH_asp = ham.dH_asp 

95 

96 for a, dP_niv in dP_aniv.items(): 

97 dP_niv = dP_niv.conj() 

98 dO_ii = wfs.setups[a].dO_ii 

99 P_ni = P_ani[a][:n_occ] 

100 dS_nkv = np.einsum('niv,ij,kj->nkv', dP_niv, dO_ii, P_ni) 

101 dS_nkv = \ 

102 (dS_nkv + np.transpose(dS_nkv, (1, 0, 2)).conj()) 

103 # '-' becuase there is extra minus in 'pt derivative' 

104 F_wfs_av[a] -= np.einsum('kn,nkv->v', lamb, dS_nkv).real 

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

106 dh_nv = np.einsum('niv,ij,nj->nv', dP_niv, dH_ii, P_ni) 

107 # '+' becuase there is extra minus in 'pt derivative' 

108 F_wfs_av[a] += np.einsum('n,nv->v', kpt.f_n[:n_occ], 

109 2.0 * dh_nv.real) 

110 

111 if 'SIC' in esolv.odd.name: 

112 esolv.odd.get_odd_corrections_to_forces(F_wfs_av, 

113 wfs, 

114 kpt) 

115 

116 wfs.kd.comm.sum(F_wfs_av, 0) 

117 wfs.gd.comm.sum(F_wfs_av, 0) 

118 

119 F_ham_av = np.zeros((natoms, 3)) 

120 ham.calculate_forces(dens, F_ham_av) 

121 F_av = F_ham_av + F_wfs_av 

122 wfs.world.broadcast(F_av, 0) 

123 

124 F_av = wfs.kd.symmetry.symmetrize_forces(F_av) 

125 

126 if log: 

127 log('\nForces in eV/Ang:') 

128 c = Hartree / Bohr 

129 for a, setup in enumerate(wfs.setups): 

130 log('%3d %-2s %10.5f %10.5f %10.5f' % 

131 ((a, setup.symbol) + tuple(F_av[a] * c))) 

132 log() 

133 

134 return F_av