Coverage for local_installation/dynasor/post_processing/spherical_average.py: 91%

93 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2025-04-16 06:13 +0000

1import numpy as np 

2 

3from copy import deepcopy 

4from dynasor.logging_tools import logger 

5from dynasor.sample import Sample 

6from numpy.typing import NDArray 

7from scipy.stats import norm 

8 

9 

10def get_spherically_averaged_sample_smearing( 

11 sample: Sample, q_norms: NDArray[float], q_width: float, sum: bool = False) -> Sample: 

12 r""" 

13 Compute a spherical average over q-points for all the correlation functions in :attr:`sample`. 

14 

15 In the gaussian average method each q-point contributes to the function value at 

16 given :math:`\vec{q}` with a weight determined by a gaussian function. For example 

17 

18 .. math:: 

19 

20 F(q) = \sum_i w(\boldsymbol{q}_i, q) F(\boldsymbol{q}_i) 

21 

22 where 

23 

24 .. math:: 

25 

26 w(\boldsymbol{q}_i, q) \propto \exp{\left [ -\frac{1}{2} \left ( \frac{|\boldsymbol{q}_i| 

27 - q}{q_{width}} \right)^2 \right ]} 

28 

29 and 

30 

31 .. math:: 

32 

33 \sum_i w(\boldsymbol{q}_i, q) = 1.0 

34 

35 This corresponds to a gaussian smearing or convolution. 

36 The input parameters are :attr:`q_norms`, setting to the values of :math:`|\vec{q}|`, 

37 for which the function is evaluated and :attr:`q_width` specifying the 

38 standard deviation of the gaussian smearing. 

39 

40 Parameters 

41 ---------- 

42 sample 

43 Input sample. 

44 q_norms 

45 Values of :math:`|\vec{q}|` at which to evaluate the correlation functions. 

46 q_width 

47 Standard deviation of the gaussian smearing. 

48 sum 

49 Whether to average or sum the sample in each bin. 

50 """ 

51 if not isinstance(sample, Sample): 

52 raise ValueError('Input sample is not a Sample object.') 

53 

54 # get q-points 

55 q_points = sample.q_points 

56 if q_points.shape[1] != 3: 

57 raise ValueError('q-points array has the wrong shape.') 

58 

59 # setup new input dicts for new Sample, remove q_points, add q_norms 

60 meta_data = deepcopy(sample.meta_data) 

61 data_dict = dict() 

62 for key in sample.dimensions: 

63 if key == 'q_points': 

64 continue 

65 data_dict[key] = sample[key] 

66 

67 for key in sample.available_correlation_functions: 

68 Z = getattr(sample, key) 

69 if sum: 69 ↛ 70line 69 didn't jump to line 70, because the condition on line 69 was never true

70 averaged_data = _get_gaussian_sum(q_points, Z, q_norms, q_width) 

71 else: 

72 averaged_data = _get_gaussian_average(q_points, Z, q_norms, q_width) 

73 data_dict[key] = averaged_data 

74 data_dict['q_norms'] = q_norms 

75 

76 return sample.__class__(data_dict, **meta_data) 

77 

78 

79def get_spherically_averaged_sample_binned( 

80 sample: Sample, num_q_bins: int, sum: bool = False) -> Sample: 

81 r""" 

82 Compute a spherical average over q-points for all the correlation functions in `:attr:`sample`. 

83 

84 Here, a q-binning method is used to conduct the spherical average, meaning all q-points are 

85 placed into spherical bins (shells). 

86 The corresponding function is calculated as the average of all q-points in a bin. 

87 If a q-bin does not contain any q-points, then its value is set to ``np.nan``. 

88 The q_min and q_max are determined from min/max of ``|q_points|``, and will determine 

89 the q-bin range. 

90 These will be set as bin-centers for the first and last bins repsectivley. 

91 The input parameter is the number of q-bins to use :attr:`num_q_bins`. 

92 

93 Parameters 

94 ---------- 

95 sample 

96 Input sample. 

97 num_q_bins 

98 Number of q-bins to use. 

99 sum 

100 Whether to average or sum the sample in each bin. 

101 """ 

102 

103 if not isinstance(sample, Sample): 

104 raise ValueError('input sample is not a Sample object.') 

105 

106 # get q-points 

107 q_points = sample.q_points 

108 if q_points.shape[1] != 3: 

109 raise ValueError('q-points array has wrong shape.') 

110 

111 # setup new input dicts for new Sample, remove q_points, add q_norms 

112 meta_data = deepcopy(sample.meta_data) 

113 data_dict = dict() 

114 for key in sample.dimensions: 

115 if key == 'q_points': 

116 continue 

117 data_dict[key] = sample[key] 

118 

119 # compute spherical average for each correlation function 

120 for key in sample.available_correlation_functions: 

121 Z = getattr(sample, key) 

122 q_bincenters, bin_counts, averaged_data = _get_bin_average(q_points, Z, num_q_bins, sum) 

123 data_dict[key] = averaged_data 

124 data_dict['q_norms'] = q_bincenters 

125 

126 return sample.__class__(data_dict, **meta_data) 

127 

128 

129def _get_gaussian_average( 

130 q_points: np.ndarray, Z: np.ndarray, q_norms: np.ndarray, q_width: float): 

131 

132 q_norms_sample = np.linalg.norm(q_points, axis=1) 

133 Z_average = [] 

134 for q in q_norms: 

135 weights = _gaussian(q_norms_sample, x0=q, sigma=q_width).reshape(-1, 1) 

136 norm = np.sum(weights) 

137 if norm != 0: 137 ↛ 139line 137 didn't jump to line 139, because the condition on line 137 was never false

138 weights = weights / norm 

139 Z_average.append(np.sum(weights * Z, axis=0)) 

140 return np.array(Z_average) 

141 

142 

143def _get_gaussian_sum( 

144 q_points: np.ndarray, Z: np.ndarray, q_norms: np.ndarray, q_width: float): 

145 

146 q_norms_sample = np.linalg.norm(q_points, axis=1) 

147 Z_average = [] 

148 for q in q_norms: 

149 weights = _gaussian(q_norms_sample, x0=q, sigma=q_width).reshape(-1, 1) 

150 Z_average.append(np.sum(weights * Z, axis=0)) 

151 return np.array(Z_average) 

152 

153 

154def _gaussian(x, x0, sigma): 

155 dist = norm(loc=x0, scale=sigma) 

156 return dist.pdf(x) 

157 

158 

159def _get_bin_average(q_points: np.ndarray, data: np.ndarray, num_q_bins: int, sum: bool = False): 

160 """ 

161 Compute a spherical average over q-points for the data using q-bins. 

162 

163 If a q-bin does not contain any q-points, then a np.nan is inserted. 

164 

165 The q_min and q_min are determined from min/max of |q_points|, and will determine the bin-range. 

166 These will set as bin-centers for the first and last bins repsectivley. 

167 

168 Parameters 

169 ---------- 

170 q_points 

171 array of q-points shape ``(Nq, 3)`` 

172 data 

173 data-array of shape ``(Nq, N)``, shape cannot be ``(Nq, )`` 

174 num_q_bins 

175 number of radial q-point bins to use 

176 sum 

177 whether or not to sum the data in each bin 

178 

179 Returns 

180 ------- 

181 q 

182 array of |q| bins of shape ``(num_q_bins, )`` 

183 data_averaged 

184 averaged data-array of shape `` 

185 """ 

186 N_qpoints = q_points.shape[0] 

187 N_t = data.shape[1] 

188 assert q_points.shape[1] == 3 

189 assert data.shape[0] == N_qpoints 

190 

191 # q-norms 

192 q_norms = np.linalg.norm(q_points, axis=1) 

193 assert q_norms.shape == (N_qpoints,) 

194 

195 # setup bins 

196 q_max = np.max(q_norms) 

197 q_min = np.min(q_norms) 

198 delta_x = (q_max - q_min) / (num_q_bins - 1) 

199 q_range = (q_min - delta_x / 2, q_max + delta_x / 2) 

200 bin_counts, edges = np.histogram(q_norms, bins=num_q_bins, range=q_range) 

201 q_bincenters = 0.5 * (edges[1:] + edges[:-1]) 

202 

203 # calculate average for each bin 

204 averaged_data = np.zeros((num_q_bins, N_t)) 

205 for bin_index in range(num_q_bins): 

206 # find q-indices that belong to this bin 

207 bin_min = edges[bin_index] 

208 bin_max = edges[bin_index + 1] 

209 bin_count = bin_counts[bin_index] 

210 q_indices = np.where(np.logical_and(q_norms >= bin_min, q_norms < bin_max))[0] 

211 assert len(q_indices) == bin_count 

212 logger.debug(f'bin {bin_index} contains {bin_count} q-points') 

213 

214 # average over q-indices, if no indices then np.nan 

215 if bin_count == 0: 

216 logger.warning(f'No q-points for bin {bin_index}') 

217 data_bin = np.array([np.nan for _ in range(N_t)]) 

218 else: 

219 if sum: 

220 data_bin = data[q_indices, :].sum(axis=0) 

221 else: 

222 data_bin = data[q_indices, :].mean(axis=0) 

223 averaged_data[bin_index, :] = data_bin 

224 

225 return q_bincenters, bin_counts, averaged_data