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
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, sum: bool = False) -> 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 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.')
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.')
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]
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
76 return sample.__class__(data_dict, **meta_data)
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`.
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`.
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 """
103 if not isinstance(sample, Sample):
104 raise ValueError('input sample is not a Sample object.')
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.')
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]
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
126 return sample.__class__(data_dict, **meta_data)
129def _get_gaussian_average(
130 q_points: np.ndarray, Z: np.ndarray, q_norms: np.ndarray, q_width: float):
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)
143def _get_gaussian_sum(
144 q_points: np.ndarray, Z: np.ndarray, q_norms: np.ndarray, q_width: float):
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)
154def _gaussian(x, x0, sigma):
155 dist = norm(loc=x0, scale=sigma)
156 return dist.pdf(x)
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.
163 If a q-bin does not contain any q-points, then a np.nan is inserted.
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.
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
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
191 # q-norms
192 q_norms = np.linalg.norm(q_points, axis=1)
193 assert q_norms.shape == (N_qpoints,)
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])
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')
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
225 return q_bincenters, bin_counts, averaged_data