Coverage for dynasor/post_processing/atomic_weighting.py: 99%
77 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-03 19:46 +0000
1from copy import deepcopy
2from typing import Optional
3from warnings import warn
4import numpy as np
5from dynasor.post_processing.weights import Weights
6from dynasor.sample import Sample, StaticSample, DynamicSample
7from numpy.typing import NDArray
10def get_weighted_sample(sample: Sample,
11 weights: Weights,
12 atom_type_map: Optional[dict[str, str]] = None) -> Sample:
13 r"""
14 Weights correlation functions with atomic weighting factors
16 The weighting of a partial dynamic structure factor
17 :math:`S_\mathrm{AB}(\boldsymbol{q}, \omega)`
18 for atom types :math:`A` and :math:`B` is carried out as
20 .. math::
22 S_\mathrm{AB}(\boldsymbol{q}, \omega)
23 = f_\mathrm{A}(\boldsymbol{q}) f_\mathrm{B}(\boldsymbol{q})
24 S_\mathrm{AB}(\boldsymbol{q}, \omega)
26 :math:`f_\mathrm{A}(\boldsymbol{q})` and :math:`f_\mathrm{B}(\boldsymbol{q})`
27 are atom-type and :math:`\boldsymbol{q}`-point dependent weights.
29 If sample has incoherent correlation functions, but :attr:`weights` does not contain
30 information on how to weight the incoherent part, then it will be dropped from the
31 returned :class:`Sample <dynasor.sample.Sample>` object (and analogously for current
32 correlation functions).
34 Parameters
35 ----------
36 sample
37 Input sample to be weighted.
38 weights
39 Object containing the weights :math:`f_\mathrm{X}(\boldsymbol{q})`.
40 atom_type_map
41 Map between the atom types in the :class:`Sample <dynasor.sample.Sample>` and the
42 ones used in the :class:`Weights` object, e.g., Ba → Ba(2+).
44 Returns
45 -------
46 A :class:`Sample <dynasor.sample.Sample>` instance with the weighted partial and
47 total structure factors.
48 """
50 # check input arguments
51 if not isinstance(sample, (StaticSample, DynamicSample)):
52 raise TypeError('sample must be a StaticSample or DynamicSample.')
53 if sample.has_incoherent and not weights.supports_incoherent:
54 warn('The Weights class does not support incoherent scattering, dropping the latter '
55 'from the weighted sample.')
57 if sample.has_currents and not weights.supports_currents:
58 warn('The Weights class does not support current correlations, dropping the latter '
59 'from the weighted sample.')
61 # setup new input dicts for new Sample
62 data_dict = dict()
63 for key in sample.dimensions:
64 data_dict[key] = sample[key]
66 # Map the atom types in the sample to the types in the weights object.
67 # Useful, for instance, when using weights for charged atomic species.
68 atom_types = [(at, at) for at in sample.atom_types]
69 if atom_type_map is not None:
70 # Fallback to use the atom type from species (`at`) if it's not mapped.
71 atom_types = [(at, atom_type_map.get(at, at)) for at in sample.atom_types]
73 # generate atomic weights for each q-point and compile to arrays
74 if 'q_norms' in sample.dimensions:
75 q_norms = sample.q_norms
76 else:
77 q_norms = np.linalg.norm(sample.q_points, axis=1)
79 weights_coh = dict()
80 for at, weight_at in atom_types:
81 weight_array = np.reshape([weights.get_weight_coh(weight_at, q) for q in q_norms], (-1, 1))
82 weights_coh[at] = weight_array
83 if sample.has_incoherent and weights.supports_incoherent:
84 weights_incoh = dict()
85 for at, weight_at in atom_types:
86 weight_array = np.reshape([
87 weights.get_weight_incoh(weight_at, q) for q in q_norms
88 ], (-1, 1))
89 weights_incoh[at] = weight_array
91 # weighting of correlation functions
92 if isinstance(sample, StaticSample):
93 data_dict_Sq = _compute_weighting_coherent(sample, 'Sq', weights_coh)
94 data_dict.update(data_dict_Sq)
95 elif isinstance(sample, DynamicSample): 95 ↛ 121line 95 didn't jump to line 121 because the condition on line 95 was always true
96 # coherent
97 Fqt_coh_dict = _compute_weighting_coherent(sample, 'Fqt_coh', weights_coh)
98 data_dict.update(Fqt_coh_dict)
99 Sqw_coh_dict = _compute_weighting_coherent(sample, 'Sqw_coh', weights_coh)
100 data_dict.update(Sqw_coh_dict)
102 # incoherent
103 if sample.has_incoherent and weights.supports_incoherent:
104 Fqt_incoh_dict = _compute_weighting_incoherent(sample, 'Fqt_incoh', weights_incoh)
105 data_dict.update(Fqt_incoh_dict)
106 Sqw_incoh_dict = _compute_weighting_incoherent(sample, 'Sqw_incoh', weights_incoh)
107 data_dict.update(Sqw_incoh_dict)
109 # currents
110 if sample.has_currents and weights.supports_currents:
111 Clqt_dict = _compute_weighting_coherent(sample, 'Clqt', weights_coh)
112 data_dict.update(Clqt_dict)
113 Clqw_dict = _compute_weighting_coherent(sample, 'Clqw', weights_coh)
114 data_dict.update(Clqw_dict)
116 Ctqt_dict = _compute_weighting_coherent(sample, 'Ctqt', weights_coh)
117 data_dict.update(Ctqt_dict)
118 Ctqw_dict = _compute_weighting_coherent(sample, 'Ctqw', weights_coh)
119 data_dict.update(Ctqw_dict)
121 new_sample = sample.__class__(
122 data_dict,
123 simulation_data=deepcopy(sample.simulation_data),
124 history=deepcopy(sample.history))
125 new_sample._append_history(
126 'get_weighted_sample',
127 dict(
128 atom_type_map=atom_type_map,
129 weights_class=weights.__class__.__name__,
130 weights_parameters=weights.parameters.to_dict(),
131 ))
133 return new_sample
136def _compute_weighting_coherent(
137 sample: Sample,
138 name: str,
139 weight_dict: dict,
140) -> dict[str, NDArray[float]]:
141 """
142 Helper function for weighting and summing partial coherent correlation functions.
143 """
144 data_dict = dict()
145 total = np.zeros(sample[name].shape)
146 for s1, s2 in sample.pairs:
147 key_pair = f'{name}_{s1}_{s2}'
148 partial = np.real(np.conjugate(weight_dict[s1]) * weight_dict[s2]) * sample[key_pair]
149 data_dict[key_pair] = partial
150 total += partial
151 data_dict[name] = total
152 return data_dict
155def _compute_weighting_incoherent(
156 sample: Sample,
157 name: str,
158 weight_dict: dict,
159) -> dict[str, NDArray[float]]:
160 """
161 Helper function for weighting and summing partial incoherent correlation functions.
162 """
163 data_dict = dict()
164 total = np.zeros(sample[name].shape)
165 for s1 in sample.atom_types:
166 key = f'{name}_{s1}'
167 partial = np.real(np.conjugate(weight_dict[s1]) * weight_dict[s1]) * sample[key]
168 data_dict[key] = partial
169 total += partial
170 data_dict[name] = total
171 return data_dict