Coverage for gpaw/pw/density.py: 96%

1from math import pi 

2import numpy as np 

3 

4from gpaw.density import Density 

5from gpaw.pw.descriptor import PWDescriptor, PWMapping 

6from gpaw.pw.lfc import PWLFC 

7 

8 

9class PseudoCoreKineticEnergyDensityLFC(PWLFC): 

10 def add(self, tauct_R): 

11 tauct_R += self.pd.ifft(1.0 / self.pd.gd.dv * 

12 self.expand().sum(1).view(complex)) 

13 

14 def derivative(self, dedtaut_R, dF_aiv): 

15 PWLFC.derivative(self, self.pd.fft(dedtaut_R), dF_aiv) 

16 

17 

18class ReciprocalSpaceDensity(Density): 

19 def __init__(self, ecut, 

20 gd, finegd, nspins, collinear, charge, redistributor, 

21 background_charge=None): 

22 Density.__init__(self, gd, finegd, nspins, collinear, charge, 

23 redistributor=redistributor, 

24 background_charge=background_charge) 

25 ecut0 = 0.5 * pi**2 / (gd.h_cv**2).sum(1).max() 

26 ecut = min(ecut, ecut0) 

27 self.pd2 = PWDescriptor(ecut, gd) 

28 self.pd3 = PWDescriptor(4 * ecut, finegd) 

29 

30 self.map23 = PWMapping(self.pd2, self.pd3) 

31 

32 self.nct_q = None 

33 self.nt_Q = None 

34 self.rhot_q = None 

35 

36 def initialize(self, setups, timer, magmom_av, hund): 

37 Density.initialize(self, setups, timer, magmom_av, hund) 

38 

39 spline_aj = [] 

40 for setup in setups: 

41 if setup.nct is None: 

42 spline_aj.append([]) 

43 else: 

44 spline_aj.append([setup.nct]) 

45 self.nct = PWLFC(spline_aj, self.pd2) 

46 

47 self.ghat = PWLFC([setup.ghat_l for setup in setups], self.pd3, 

48 ) # blocksize=256, comm=self.xc_redistributor.comm) 

49 

50 def set_positions(self, spos_ac, atom_partition): 

51 Density.set_positions(self, spos_ac, atom_partition) 

52 self.nct_q = self.pd2.zeros() 

53 self.nct.add(self.nct_q, 1.0 / self.nspins) 

54 self.nct_G = self.pd2.ifft(self.nct_q) 

55 

56 def interpolate_pseudo_density(self, comp_charge=None): 

57 """Interpolate pseudo density to fine grid.""" 

58 if comp_charge is None: 

59 comp_charge, _Q_aL = self.calculate_multipole_moments() 

60 

61 if self.nt_xg is None: 

62 self.nt_xg = self.finegd.empty(self.ncomponents) 

63 self.nt_sg = self.nt_xg[:self.nspins] 

64 self.nt_vg = self.nt_xg[self.nspins:] 

65 self.nt_Q = self.pd2.empty() 

66 

67 self.nt_Q[:] = 0.0 

68 

69 x = 0 

70 for nt_G, nt_g in zip(self.nt_xG, self.nt_xg): 

71 nt_g[:], nt_Q = self.pd2.interpolate(nt_G, self.pd3) 

72 if x < self.nspins: 

73 self.nt_Q += nt_Q 

74 x += 1 

75 

76 def interpolate(self, in_xR, out_xR=None): 

77 """Interpolate array(s).""" 

78 if out_xR is None: 

79 out_xR = self.finegd.empty(in_xR.shape[:-3]) 

80 

81 a_xR = in_xR.reshape((-1,) + in_xR.shape[-3:]) 

82 b_xR = out_xR.reshape((-1,) + out_xR.shape[-3:]) 

83 

84 for in_R, out_R in zip(a_xR, b_xR): 

85 out_R[:] = self.pd2.interpolate(in_R, self.pd3)[0] 

86 

87 return out_xR 

88 

89 distribute_and_interpolate = interpolate 

90 

91 def calculate_pseudo_charge(self): 

92 self.rhot_q = self.pd3.zeros() 

93 Q_aL = self.Q.calculate(self.D_asp) 

94 self.ghat.add(self.rhot_q, Q_aL) 

95 self.map23.add_to2(self.rhot_q, self.nt_Q) 

96 self.background_charge.add_fourier_space_charge_to(self.pd3, 

97 self.rhot_q) 

98 

99 def get_pseudo_core_kinetic_energy_density_lfc(self): 

100 return PseudoCoreKineticEnergyDensityLFC( 

101 [[setup.tauct] for setup in self.setups], self.pd2) 

102 

103 def calculate_dipole_moment(self): 

104 pd = self.pd3 

105 N_c = pd.tmp_Q.shape 

106 

107 m0_q, m1_q, m2_q = (i_G == 0 

108 for i_G in np.unravel_index(pd.Q_qG[0], N_c)) 

109 rhot_q = self.pd3.gather(self.rhot_q) 

110 if pd.comm.rank == 0: 

111 irhot_q = rhot_q.imag 

112 rhot_cs = [irhot_q[m1_q & m2_q], 

113 irhot_q[m0_q & m2_q], 

114 irhot_q[m0_q & m1_q]] 

115 d_c = [np.dot(rhot_s[1:], 1.0 / np.arange(1, len(rhot_s))) 

116 for rhot_s in rhot_cs] 

117 d_v = -np.dot(d_c, pd.gd.cell_cv) / pi * pd.gd.dv 

118 else: 

119 d_v = np.empty(3) 

120 pd.comm.broadcast(d_v, 0) 

121 return d_v