Coverage for gpaw/response/frequencies.py: 96%

1from __future__ import annotations 

2 

3import numbers 

4from typing import Any 

5 

6import numpy as np 

7from ase.units import Ha 

8from gpaw.typing import ArrayLike1D 

9 

10 

11class FrequencyDescriptor: 

12 def __init__(self, omega_w: ArrayLike1D): 

13 """Frequency grid descriptor. 

14 

15 Parameters 

16 ---------- 

17 omega_w: 

18 Frequency grid in Hartree units. 

19 """ 

20 self.omega_w = np.asarray(omega_w).copy() 

21 

22 def __len__(self): 

23 return len(self.omega_w) 

24 

25 def __repr__(self): 

26 emin = self.omega_w[0] * Ha 

27 emax = self.omega_w[-1] * Ha 

28 return (f'{self.__class__.__name__}' 

29 f'(from {emin:.3f} to {emax:.3f} eV, {len(self)} points)') 

30 

31 @staticmethod 

32 def from_array_or_dict(input: dict[str, Any] | ArrayLike1D 

33 ) -> FrequencyDescriptor: 

34 """Create frequency-grid descriptor. 

35 

36 In case *input* is a list on frequencies (in eV) a 

37 :class:`FrequencyGridDescriptor` instance is returned. 

38 Othervise a :class:`NonLinearFrequencyDescriptor` instance is 

39 returned. 

40 

41 >>> from ase.units import Ha 

42 >>> params = dict(type='nonlinear', 

43 ... domega0=0.1, 

44 ... omega2=10, 

45 ... omegamax=50) 

46 >>> wd = FrequencyDescriptor.from_array_or_dict(params) 

47 >>> wd.omega_w[0:2] * Ha 

48 array([0. , 0.10041594]) 

49 """ 

50 if isinstance(input, dict): 

51 assert input['type'] == 'nonlinear' 

52 domega0 = input.get('domega0') 

53 omega2 = input.get('omega2') 

54 omegamax = input['omegamax'] 

55 return NonLinearFrequencyDescriptor( 

56 (0.1 if domega0 is None else domega0) / Ha, 

57 (10.0 if omega2 is None else omega2) / Ha, 

58 omegamax / Ha) 

59 return FrequencyGridDescriptor(np.asarray(input) / Ha) 

60 

61 

62class ComplexFrequencyDescriptor: 

63 

64 def __init__(self, hz_z: ArrayLike1D): 

65 """Construct the complex frequency descriptor. 

66 

67 Parameters 

68 ---------- 

69 hz_z: 

70 Array of complex frequencies (in units of Hartree) 

71 """ 

72 # Use a copy of the input array 

73 hz_z = np.array(hz_z) 

74 assert hz_z.dtype == complex 

75 

76 self.hz_z = hz_z 

77 

78 def __len__(self): 

79 return len(self.hz_z) 

80 

81 def almost_eq(self, zd): 

82 if len(zd) != len(self): 

83 return False 

84 return np.allclose(self.hz_z, zd.hz_z) 

85 

86 @staticmethod 

87 def from_array(frequencies: ArrayLike1D): 

88 """Create a ComplexFrequencyDescriptor from frequencies in eV.""" 

89 return ComplexFrequencyDescriptor(np.asarray(frequencies) / Ha) 

90 

91 @property 

92 def upper_half_plane(self): 

93 """All frequencies reside in the upper half complex frequency plane?""" 

94 return np.all(self.hz_z.imag > 0.) 

95 

96 @property 

97 def horizontal_contour(self): 

98 """Do all frequency point lie on a horizontal contour?""" 

99 return np.ptp(self.hz_z.imag) < 1.e-6 

100 

101 @property 

102 def omega_w(self): 

103 """Real part of the frequencies.""" 

104 assert self.horizontal_contour, \ 

105 'It only makes sense to index the frequencies by their real part '\ 

106 'if they reside on a horizontal contour.' 

107 return self.hz_z.real 

108 

109 

110class FrequencyGridDescriptor(FrequencyDescriptor): 

111 

112 def get_index_range(self, lim1_m, lim2_m): 

113 """Get index range. """ 

114 

115 i0_m = np.zeros(len(lim1_m), int) 

116 i1_m = np.zeros(len(lim2_m), int) 

117 

118 for m, (lim1, lim2) in enumerate(zip(lim1_m, lim2_m)): 

119 i_x = np.logical_and(lim1 <= self.omega_w, 

120 lim2 >= self.omega_w) 

121 if i_x.any(): 

122 inds = np.argwhere(i_x) 

123 i0_m[m] = inds.min() 

124 i1_m[m] = inds.max() + 1 

125 

126 return i0_m, i1_m 

127 

128 

129class NonLinearFrequencyDescriptor(FrequencyDescriptor): 

130 def __init__(self, 

131 domega0: float, 

132 omega2: float, 

133 omegamax: float): 

134 """Non-linear frequency grid. 

135 

136 Units are Hartree. See :ref:`frequency grid`. 

137 

138 Parameters 

139 ---------- 

140 domega0: 

141 Frequency grid spacing for non-linear frequency grid at omega = 0. 

142 omega2: 

143 Frequency at which the non-linear frequency grid has doubled 

144 the spacing. 

145 omegamax: 

146 The upper frequency bound for the non-linear frequency grid. 

147 """ 

148 beta = (2**0.5 - 1) * domega0 / omega2 

149 wmax = int(omegamax / (domega0 + beta * omegamax)) 

150 w = np.arange(wmax + 2) # + 2 is for buffer 

151 omega_w = w * domega0 / (1 - beta * w) 

152 

153 super().__init__(omega_w) 

154 

155 self.domega0 = domega0 

156 self.omega2 = omega2 

157 self.omegamax = omegamax 

158 self.omegamin = 0 

159 

160 self.beta = beta 

161 self.wmax = wmax 

162 self.omega_w = omega_w 

163 self.wmax = wmax 

164 

165 def get_floor_index(self, o_m, safe=True): 

166 """Get closest index rounding down.""" 

167 beta = self.beta 

168 w_m = (o_m / (self.domega0 + beta * o_m)).astype(int) 

169 if safe: 

170 if isinstance(w_m, np.ndarray): 

171 w_m[w_m >= self.wmax] = self.wmax - 1 

172 elif isinstance(w_m, numbers.Integral): 

173 if w_m >= self.wmax: 

174 w_m = self.wmax - 1 

175 else: 

176 raise TypeError 

177 return w_m 

178 

179 def get_index_range(self, omega1_m, omega2_m): 

180 omega1_m = omega1_m.copy() 

181 omega2_m = omega2_m.copy() 

182 omega1_m[omega1_m < 0] = 0 

183 omega2_m[omega2_m < 0] = 0 

184 w1_m = self.get_floor_index(omega1_m) 

185 w2_m = self.get_floor_index(omega2_m) 

186 o1_m = self.omega_w[w1_m] 

187 o2_m = self.omega_w[w2_m] 

188 w1_m[o1_m < omega1_m] += 1 

189 w2_m[o2_m < omega2_m] += 1 

190 return w1_m, w2_m