Coverage for gpaw/response/qpd.py: 79%

1from math import pi 

2import numpy as np 

3 

4from gpaw.kpt_descriptor import KPointDescriptor 

5from gpaw.pw.descriptor import PWDescriptor 

6import gpaw.fftw as fftw 

7 

8 

9class SingleQPWDescriptor(PWDescriptor): 

10 @staticmethod 

11 def from_q(q_c, ecut, gd, gammacentered=False): 

12 """Construct a plane wave descriptor for q_c with a given cutoff.""" 

13 qd = KPointDescriptor([q_c]) 

14 if not isinstance(ecut, dict): 

15 return SingleQPWDescriptor(ecut, gd, complex, qd, 

16 gammacentered=gammacentered) 

17 elif ecut['class'] is SingleCylQPWDescriptor: 

18 return SingleCylQPWDescriptor(gd=gd, 

19 kd=qd, 

20 gammacentered=gammacentered, 

21 dtype=complex, 

22 **ecut['kwargs']) 

23 else: 

24 raise NotImplementedError( 

25 f'Unrecognized QPW class: {ecut["class"]}') 

26 

27 @property 

28 def q_c(self): 

29 return self.kd.bzk_kc[0] 

30 

31 @property 

32 def optical_limit(self): 

33 return np.allclose(self.q_c, 0.0) 

34 

35 def copy(self): 

36 return self.copy_with() 

37 

38 def copy_with(self, ecut=None, gd=None, gammacentered=None): 

39 if ecut is None: 

40 ecut = self.ecut 

41 if gd is None: 

42 gd = self.gd 

43 if gammacentered is None: 

44 gammacentered = self.gammacentered 

45 

46 return SingleQPWDescriptor.from_q( 

47 self.q_c, ecut, gd, gammacentered=gammacentered) 

48 

49 

50class SingleCylQPWDescriptor(SingleQPWDescriptor): 

51 def __init__(self, ecut_xy, ecut_z, gd, dtype=None, kd=None, 

52 fftwflags=fftw.MEASURE, gammacentered=False): 

53 

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

55 if ecut_xy is None: 

56 ecut_xy = 0.9999 * ecut0 

57 else: 

58 assert ecut_xy <= ecut0 

59 

60 if ecut_z is None: 

61 ecut_z = 0.9999 * ecut0 

62 else: 

63 assert ecut_z <= ecut0 

64 

65 self.ecut_z = ecut_z 

66 self.ecut_xy = ecut_xy 

67 

68 super().__init__(ecut_xy, gd, dtype, kd, 

69 fftwflags, gammacentered) 

70 

71 def setup_pw_grid(self, i_Qc, Q_Q): 

72 ng_q = [] 

73 Q_qG = [] 

74 G2_qG = [] 

75 for q, K_v in enumerate(self.K_qv): 

76 G2_Q = ((self.G_Qv + K_v)**2).sum(axis=1) 

77 if self.gammacentered: 

78 mask_Q = ((self.G_Qv[:, 0:2]**2).sum(axis=1) 

79 <= 2 * self.ecut_xy) \ 

80 & ((self.G_Qv[:, 2]**2) <= 2 * self.ecut_z) 

81 else: 

82 G3_Q = ((self.G_Qv[:, 0:2] + K_v[0:2])**2).sum(axis=1) 

83 mask_Q = (G3_Q <= (2 * self.ecut_xy)) \ 

84 & ((self.G_Qv[:, 2]**2) <= (2 * self.ecut_z)) 

85 

86 if self.dtype == float: 

87 mask_Q &= ((i_Qc[:, 2] > 0) | 

88 (i_Qc[:, 1] > 0) | 

89 ((i_Qc[:, 0] >= 0) & (i_Qc[:, 1] == 0))) 

90 Q_G = Q_Q[mask_Q] 

91 Q_qG.append(Q_G) 

92 G2_qG.append(G2_Q[Q_G]) 

93 ng = len(Q_G) 

94 ng_q.append(ng) 

95 

96 return ng_q, Q_qG, G2_qG 

97 

98 def copy_with(self, ecut=None, gd=None, gammacentered=None): 

99 if ecut is None: 

100 ecut = self.ecut_xy 

101 if gd is None: 

102 gd = self.gd 

103 if gammacentered is None: 

104 gammacentered = self.gammacentered 

105 

106 return SingleCylQPWDescriptor.from_q( 

107 self.q_c, ecut_xy=ecut, ecut_z=self.ecut_z, 

108 gd=gd, gammacentered=gammacentered)