Coverage for gpaw/stress.py: 97%

68 statements  

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

1import numpy as np 

2import ase.units as units 

3 

4from gpaw.utilities import unpack_hermitian 

5from gpaw.wavefunctions.pw import PWWaveFunctions 

6 

7 

8def calculate_stress(calc): 

9 wfs = calc.wfs 

10 dens = calc.density 

11 ham = calc.hamiltonian 

12 

13 if not isinstance(wfs, PWWaveFunctions): 

14 raise NotImplementedError('Calculation of stress tensor is only ' + 

15 'implemented for plane-wave mode.') 

16 if ham.xc.orbital_dependent and ham.xc.type != 'MGGA': 

17 raise NotImplementedError('Calculation of stress tensor is not ' + 

18 'implemented for orbital-dependent ' + 

19 'XC functionals such as ' + ham.xc.name) 

20 

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

22 

23 calc.timer.start('Stress tensor') 

24 

25 s_vv = wfs.get_kinetic_stress().real 

26 nt_xg = dens.nt_xg 

27 if ham.xc_redistributor is not None: 

28 nt_xg = ham.xc_redistributor.distribute(nt_xg) 

29 s_vv += ham.xc.stress_tensor_contribution(nt_xg) 

30 

31 pd = dens.pd3 

32 p_G = 4 * np.pi * dens.rhot_q 

33 G0 = 0 if pd.gd.comm.rank > 0 else 1 

34 p_G[G0:] /= pd.G2_qG[0][G0:]**2 

35 G_Gv = pd.get_reciprocal_vectors(add_q=False) 

36 for v1 in range(3): 

37 for v2 in range(3): 

38 s_vv[v1, v2] += pd.integrate(p_G, dens.rhot_q * 

39 G_Gv[:, v1] * G_Gv[:, v2]) 

40 s_vv += dens.ghat.stress_tensor_contribution(ham.vHt_q, dens.Q_aL) 

41 

42 s_vv -= np.eye(3) * ham.estress 

43 s_vv += ham.vbar.stress_tensor_contribution(dens.nt_Q) 

44 s_vv += dens.nct.stress_tensor_contribution(ham.vt_Q) 

45 if ham.xc.type == 'MGGA': 

46 nspin = ham.xc.dedtaut_sG.shape[0] 

47 vtau_sQ = dens.pd2.empty((nspin,), global_array=False) 

48 for s in range(nspin): 

49 vtau_sQ[s] = dens.pd2.fft(ham.xc.dedtaut_sG[s]) 

50 s_vv += ham.xc.tauct.stress_tensor_contribution(vtau_sQ.mean(axis=0)) 

51 

52 s0 = 0.0 

53 s0_vv = 0.0 

54 for kpt in wfs.kpt_u: 

55 a_ani = {} 

56 for a, P_ni in kpt.P_ani.items(): 

57 Pf_ni = P_ni * kpt.f_n[:, None] 

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

59 dS_ii = ham.setups[a].dO_ii 

60 a_ni = (np.dot(Pf_ni, dH_ii) - 

61 np.dot(Pf_ni * kpt.eps_n[:, None], dS_ii)) 

62 s0 += np.vdot(P_ni, a_ni) 

63 a_ani[a] = 2 * a_ni.conj() 

64 s0_vv += wfs.pt.stress_tensor_contribution(kpt.psit_nG, a_ani, 

65 q=kpt.q) 

66 s0_vv -= dens.gd.comm.sum_scalar(s0.real) * np.eye(3) 

67 s0_vv /= dens.gd.comm.size 

68 wfs.world.sum(s0_vv) 

69 s_vv += s0_vv 

70 

71 s_vv += wfs.dedepsilon * np.eye(3) 

72 

73 vol = calc.atoms.get_volume() / units.Bohr**3 

74 s_vv = 0.5 / vol * (s_vv + s_vv.T) 

75 

76 # Symmetrize: 

77 sigma_vv = np.zeros((3, 3)) 

78 cell_cv = wfs.gd.cell_cv 

79 for U_cc in wfs.kd.symmetry.op_scc: 

80 M_vv = np.dot(np.linalg.inv(cell_cv), 

81 np.dot(U_cc, cell_cv)).T 

82 sigma_vv += np.dot(np.dot(M_vv.T, s_vv), M_vv) 

83 sigma_vv /= len(wfs.kd.symmetry.op_scc) 

84 

85 # Make sure all agree on the result (redundant calculation on 

86 # different cores involving BLAS might give slightly different 

87 # results): 

88 wfs.world.broadcast(sigma_vv, 0) 

89 

90 calc.log('Stress tensor:') 

91 for sigma_v in sigma_vv: 

92 calc.log('{:13.6f}{:13.6f}{:13.6f}' 

93 .format(*(units.Ha / units.Bohr**3 * sigma_v))) 

94 

95 calc.timer.stop('Stress tensor') 

96 

97 return sigma_vv