Coverage for dynasor/post_processing/spherical_average.py: 93%
118 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-20 20:02 +0000
1from copy import deepcopy
2from typing import Optional
3import numpy as np
4from numpy.typing import NDArray
5from scipy.stats import norm
6from dynasor.logging_tools import logger
7from dynasor.sample import Sample
10def get_spherically_averaged_sample_smearing(
11 sample: Sample,
12 q_norms: NDArray[float],
13 q_width: float,
14 use_sum: Optional[bool] = False,
15 broadening: Optional[str] = 'gaussian',
16) -> Sample:
17 r"""
18 Compute a spherical average over q-points for all the correlation functions in :attr:`sample`.
20 Each q-point contributes to the function value at a given :math:`\boldsymbol{q}` with a weight
21 determined by a broadening function. For example
23 .. math::
25 F(q) = \sum_i w(\boldsymbol{q}_i, q) F(\boldsymbol{q}_i)
27 where :math:`\sum_i w(\boldsymbol{q}_i, q) = 1` (when :attr:`use_sum` is ``False``).
29 Two broadening functions are available via :attr:`broadening`:
31 **Gaussian** (``'gaussian'``):
33 .. math::
35 w(\boldsymbol{q}_i, q) \propto \exp{\left [ -\frac{1}{2} \left ( \frac{|\boldsymbol{q}_i|
36 - q}{q_{width}} \right)^2 \right ]}
38 where :attr:`q_width` is the standard deviation :math:`\sigma`.
40 **Lorentzian** (``'lorentzian'``):
42 .. math::
44 w(\boldsymbol{q}_i, q) \propto \frac{1}{\left(|\boldsymbol{q}_i| - q\right)^2
45 + q_{width}^2}
47 where :attr:`q_width` is the half-width at half-maximum :math:`\gamma`.
49 Parameters
50 ----------
51 sample
52 Input sample.
53 q_norms
54 Values of :math:`|\vec{q}|` at which to evaluate the correlation functions.
55 q_width
56 Width of the broadening function. Standard deviation :math:`\sigma` for Gaussian;
57 half-width at half-maximum :math:`\gamma` for Lorentzian.
58 use_sum
59 Whether to average or sum the sample in each bin.
60 broadening
61 Broadening function to use. Either ``'gaussian'`` (default) or ``'lorentzian'``.
62 """
63 if not isinstance(sample, Sample):
64 raise ValueError('Input sample is not a Sample object.')
66 if broadening not in ('gaussian', 'lorentzian'):
67 raise ValueError(f"broadening must be 'gaussian' or 'lorentzian', got '{broadening}'")
69 # get q-points
70 q_points = sample.q_points
71 if q_points.shape[1] != 3:
72 raise ValueError('q-points array has the wrong shape.')
74 # warn if the smearing width is finer than the q_norms grid
75 q_norms = np.asarray(q_norms)
76 if len(q_norms) > 1: 76 ↛ 85line 76 didn't jump to line 85 because the condition on line 76 was always true
77 q_spacing = np.ptp(q_norms) / (len(q_norms) - 1)
78 if q_width < q_spacing:
79 logger.warning(
80 f'q_width ({q_width}) is smaller than the q_norms spacing ({q_spacing:.3g}); '
81 'the spherical smearing may be under-resolved (spiky). Consider increasing '
82 'q_width or using a coarser q_norms grid.')
84 # set up new input dicts for new Sample, remove q_points, add q_norms
85 data_dict = dict()
86 for key in sample.dimensions:
87 if key == 'q_points':
88 continue
89 data_dict[key] = sample[key]
91 if broadening == 'gaussian':
92 get_average = _get_gaussian_average
93 get_sum = _get_gaussian_sum
94 else:
95 get_average = _get_lorentzian_average
96 get_sum = _get_lorentzian_sum
98 for key in sample.available_correlation_functions:
99 Z = getattr(sample, key)
100 if use_sum: 100 ↛ 101line 100 didn't jump to line 101 because the condition on line 100 was never true
101 averaged_data = get_sum(q_points, Z, q_norms, q_width)
102 else:
103 averaged_data = get_average(q_points, Z, q_norms, q_width)
104 data_dict[key] = averaged_data
105 data_dict['q_norms'] = q_norms
107 # compose new object
108 new_sample = sample.__class__(
109 data_dict,
110 simulation_data=deepcopy(sample.simulation_data),
111 history=deepcopy(sample.history))
112 new_sample._append_history(
113 'get_spherically_averaged_sample_smearing',
114 dict(
115 q_width=q_width,
116 use_sum=use_sum,
117 broadening=broadening,
118 ))
120 return new_sample
123def get_spherically_averaged_sample_binned(
124 sample: Sample,
125 num_q_bins: int,
126 use_sum: Optional[bool] = False,
127) -> Sample:
128 r"""
129 Compute a spherical average over q-points for all the correlation functions in :attr:`sample`.
131 Here, a q-binning method is used to conduct the spherical average, meaning all q-points are
132 placed into spherical bins (shells).
133 The corresponding function is calculated as the average of all q-points in a bin.
134 If a q-bin does not contain any q-points, then its value is set to `np.nan`.
135 The boundaries of the range, `q_min` and `q_max`, are taken as the minimum and maximum,
136 respectively, of `|q_points|`.
137 These will be set as bin centers for the first and last bins, respectively.
138 The input parameter is the number of q-bins to use :attr:`num_q_bins`.
140 Parameters
141 ----------
142 sample
143 Input sample.
144 num_q_bins
145 Number of q-bins to use.
146 use_sum
147 Whether to average or sum the sample in each bin.
148 """
150 if not isinstance(sample, Sample):
151 raise ValueError('Input sample is not a Sample object.')
153 # get q-points
154 q_points = sample.q_points
155 if q_points.shape[1] != 3:
156 raise ValueError('q-points array has wrong shape.')
158 # set up new input dicts for new Sample, remove q_points, add q_norms
159 data_dict = dict()
160 for key in sample.dimensions:
161 if key == 'q_points':
162 continue
163 data_dict[key] = sample[key]
165 # compute spherical average for each correlation function
166 for key in sample.available_correlation_functions:
167 Z = getattr(sample, key)
168 q_bincenters, bin_counts, averaged_data = _get_bin_average(q_points, Z, num_q_bins, use_sum)
169 data_dict[key] = averaged_data
170 data_dict['q_norms'] = q_bincenters
172 # compose new sample
173 new_sample = sample.__class__(
174 data_dict,
175 simulation_data=deepcopy(sample.simulation_data),
176 history=deepcopy(sample.history))
177 new_sample._append_history(
178 'get_spherically_averaged_sample_binned',
179 dict(
180 num_q_bins=num_q_bins,
181 use_sum=use_sum,
182 ))
184 return new_sample
187def _get_gaussian_average(
188 q_points: NDArray[float],
189 Z: NDArray[float],
190 q_norms: NDArray[float],
191 q_width: float,
192) -> NDArray[float]:
193 q_norms_sample = np.linalg.norm(q_points, axis=1)
194 # Subtract the per-target max exponent (log-sum-exp trick): mathematically
195 # identical after normalization but avoids underflowing all weights to zero.
196 diff = q_norms[:, None] - q_norms_sample[None, :]
197 exponent = -0.5 * (diff / q_width) ** 2
198 exponent -= exponent.max(axis=1, keepdims=True)
199 weights = np.exp(exponent)
200 weights /= weights.sum(axis=1, keepdims=True)
201 return weights @ Z
204def _get_gaussian_sum(
205 q_points: NDArray[float],
206 Z: NDArray[float],
207 q_norms: NDArray[float],
208 q_width: float,
209) -> NDArray[float]:
210 q_norms_sample = np.linalg.norm(q_points, axis=1)
211 diff = q_norms[:, None] - q_norms_sample[None, :]
212 weights = np.exp(-0.5 * (diff / q_width) ** 2) / (q_width * np.sqrt(2 * np.pi))
213 return weights @ Z
216def _gaussian(x: NDArray[float], x0: float, sigma: float) -> NDArray[float]:
217 dist = norm(loc=x0, scale=sigma)
218 return dist.pdf(x)
221def _get_lorentzian_average(
222 q_points: NDArray[float],
223 Z: NDArray[float],
224 q_norms: NDArray[float],
225 q_width: float,
226) -> NDArray[float]:
227 q_norms_sample = np.linalg.norm(q_points, axis=1)
228 # weights shape: (N_q_out, N_qpoints); normalization constant cancels so omit it
229 diff = q_norms[:, None] - q_norms_sample[None, :]
230 weights = 1.0 / (diff ** 2 + q_width ** 2)
231 norms = weights.sum(axis=1, keepdims=True)
232 weights /= np.where(norms != 0, norms, 1.0)
233 return weights @ Z
236def _get_lorentzian_sum(
237 q_points: NDArray[float],
238 Z: NDArray[float],
239 q_norms: NDArray[float],
240 q_width: float,
241) -> NDArray[float]:
242 q_norms_sample = np.linalg.norm(q_points, axis=1)
243 diff = q_norms[:, None] - q_norms_sample[None, :]
244 weights = q_width / (np.pi * (diff ** 2 + q_width ** 2))
245 return weights @ Z
248def _lorentzian(x: NDArray[float], x0: float, gamma: float) -> NDArray[float]:
249 return gamma / (np.pi * ((x - x0) ** 2 + gamma ** 2))
252def _get_bin_average(
253 q_points: NDArray[float],
254 data: NDArray[float],
255 num_q_bins: int,
256 use_sum: Optional[bool] = False,
257) -> tuple[NDArray[float], NDArray[float]]:
258 """
259 Compute a spherical average over q-points for the data using q-bins.
261 If a q-bin does not contain any q-points, then a np.nan is inserted.
263 q_min and q_max are determined from min/max of |q_points| and define the bin range.
264 These are set as bin centers for the first and last bins, respectively.
266 Parameters
267 ----------
268 q_points
269 Array of q-points shape ``(Nq, 3)``.
270 data
271 Array of shape ``(Nq, N)``, shape cannot be ``(Nq, )``.
272 num_q_bins
273 Number of radial q-point bins to use.
274 use_sum
275 Whether or not to sum the data in each bin.
277 Returns
278 -------
279 Tuple comprising the array of |q| bins of shape ``(num_q_bins, )``
280 and the averaged data-array.
281 """
282 N_qpoints = q_points.shape[0]
283 N_t = data.shape[1]
284 assert q_points.shape[1] == 3
285 assert data.shape[0] == N_qpoints
287 # q-norms
288 q_norms = np.linalg.norm(q_points, axis=1)
289 assert q_norms.shape == (N_qpoints,)
291 # set up bins
292 q_max = np.max(q_norms)
293 q_min = np.min(q_norms)
294 delta_x = (q_max - q_min) / (num_q_bins - 1)
295 q_range = (q_min - delta_x / 2, q_max + delta_x / 2)
296 bin_counts, edges = np.histogram(q_norms, bins=num_q_bins, range=q_range)
297 q_bincenters = 0.5 * (edges[1:] + edges[:-1])
299 # calculate average for each bin
300 averaged_data = np.zeros((num_q_bins, N_t))
301 for bin_index in range(num_q_bins):
302 # find q-indices that belong to this bin
303 bin_min = edges[bin_index]
304 bin_max = edges[bin_index + 1]
305 bin_count = bin_counts[bin_index]
306 q_indices = np.where(np.logical_and(q_norms >= bin_min, q_norms < bin_max))[0]
307 assert len(q_indices) == bin_count
308 logger.debug(f'bin {bin_index} contains {bin_count} q-points')
310 # average over q-indices, if no indices then np.nan
311 if bin_count == 0:
312 logger.warning(f'No q-points for bin {bin_index}')
313 data_bin = np.array([np.nan for _ in range(N_t)])
314 else:
315 if use_sum:
316 data_bin = data[q_indices, :].sum(axis=0)
317 else:
318 data_bin = data[q_indices, :].mean(axis=0)
319 averaged_data[bin_index, :] = data_bin
321 return q_bincenters, bin_counts, averaged_data