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