# Coverage for local_installation/dynasor/post_processing/spherical_average.py: 99%

## 82 statements

, created at 2024-08-05 09:53 +0000

1import numpy as np

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

10def get_spherically_averaged_sample_smearing(

11 sample: Sample, q_norms: NDArray[float], q_width: float) -> Sample:

12 r"""

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

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

18 .. math::

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

22 where

24 .. math::

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

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

29 and

31 .. math::

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

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.

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 """

49 if not isinstance(sample, Sample):

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

52 # get q-points

53 q_points = sample.q_points

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

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

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

58 meta_data = deepcopy(sample.meta_data)

59 data_dict = dict()

60 for key in sample.dimensions:

61 if key == 'q_points':

62 continue

63 data_dict[key] = sample[key]

65 for key in sample.available_correlation_functions:

66 Z = getattr(sample, key)

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

68 data_dict[key] = averaged_data

69 data_dict['q_norms'] = q_norms

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

74def get_spherically_averaged_sample_binned(sample: Sample, num_q_bins: int) -> Sample:

75 r"""

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

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

79 placed into spherical bins (shells).

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

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

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

83 the q-bin range.

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

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

87 Parameters

88 ----------

89 sample

90 Input sample

91 num_q_bins

92 number of q-bins to use

93 """

95 if not isinstance(sample, Sample):

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

98 # get q-points

99 q_points = sample.q_points

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

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

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

104 meta_data = deepcopy(sample.meta_data)

105 data_dict = dict()

106 for key in sample.dimensions:

107 if key == 'q_points':

108 continue

109 data_dict[key] = sample[key]

111 # compute spherical average for each correlation function

112 for key in sample.available_correlation_functions:

113 Z = getattr(sample, key)

114 q_bincenters, bin_counts, averaged_data = _get_bin_average(q_points, Z, num_q_bins)

115 data_dict[key] = averaged_data

116 data_dict['q_norms'] = q_bincenters

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

121def _get_gaussian_average(

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

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

125 Z_average = []

126 for q in q_norms:

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

128 norm = np.sum(weights)

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

130 weights = weights / norm

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

132 return np.array(Z_average)

135def _gaussian(x, x0, sigma):

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

137 return dist.pdf(x)

140def _get_bin_average(q_points: np.ndarray, data: np.ndarray, num_q_bins: int):

141 """

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

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

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

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

149 Parameters

150 ----------

151 q_points

152 array of q-points shape (Nq, 3)

153 data

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

155 num_q_bins

156 number of radial q-point bins to use

158 Returns

159 -------

160 q

161 array of |q| bins of shape (num_q_bins, )

162 data_averaged

163 averaged data-array of shape `

164 """

165 N_qpoints = q_points.shape[0]

166 N_t = data.shape[1]

167 assert q_points.shape[1] == 3

168 assert data.shape[0] == N_qpoints

170 # q-norms

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

172 assert q_norms.shape == (N_qpoints,)

174 # setup bins

175 q_max = np.max(q_norms)

176 q_min = np.min(q_norms)

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

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

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

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

182 # calculate average for each bin

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

184 for bin_index in range(num_q_bins):

185 # find q-indices that belong to this bin

186 bin_min = edges[bin_index]

187 bin_max = edges[bin_index + 1]

188 bin_count = bin_counts[bin_index]

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

190 assert len(q_indices) == bin_count

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

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

194 if bin_count == 0:

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

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

197 else:

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

199 averaged_data[bin_index, :] = data_bin

201 return q_bincenters, bin_counts, averaged_data