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

## 74 statements

, created at 2023-11-30 21:04 +0000

1import numpy as np

2from typing import Dict

3from warnings import warn

4from dynasor.post_processing.weights import Weights

5from dynasor.sample import Sample, StaticSample, DynamicSample

6from copy import deepcopy

9def get_weighted_sample(sample: Sample, weights: Weights) -> Sample:

10 r"""

11 Weights correlation functions with atomic weighting factors

13 The weighting of a partial dynamic structure factor

14 :math:S_\mathrm{AB}(\boldsymbol{q}, \omega)

15 for atom types :math:A and :math:B is carried out as

17 .. math::

19 S_\mathrm{AB}(\boldsymbol{q}, \omega)

20 = f_\mathrm{A}(\boldsymbol{q}) f_\mathrm{B}(\boldsymbol{q})

21 S_\mathrm{AB}(\boldsymbol{q}, \omega)

23 :math:f_\mathrm{A}(\boldsymbol{q}) and :math:f_\mathrm{B}(\boldsymbol{q})

24 are atom-type and :math:\boldsymbol{q}-point dependent weights.

26 If sample has incoherent correlation functions, but :attr:weights does not contain

27 information on how to weight the incoherent part, then it will be dropped from the returned

28 :attr:Sample object (and analogously for current correlation functions).

30 Parameters

31 ----------

32 sample

33 input sample to be weighted

34 weights

35 object containing the weights :math:f_\mathrm{X}(\boldsymbol{q})

37 Returns

38 -------

39 A :class:Sample instance with the weighted partial and total structure factors.

40 """

42 # check input arguments

43 if sample.has_incoherent and not weights.supports_incoherent:

44 warn('The Weights does not support incoherent scattering, dropping the latter '

45 'from the weighted sample.')

47 if sample.has_currents and not weights.supports_currents:

48 warn('The Weights does not support current correlations, dropping the latter '

49 'from the weighted sample.')

51 # setup new input dicts for new Sample

52 data_dict = dict()

53 for key in sample.dimensions:

54 data_dict[key] = sample[key]

55 meta_data = deepcopy(sample.meta_data)

57 # generate atomic weights for each q-point and compile to arrays

58 if 'q_norms' in sample.dimensions:

59 q_norms = sample.q_norms

60 else:

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

63 weights_coh = dict()

64 for at in sample.atom_types:

65 weight_array = np.reshape([weights.get_weight_coh(at, q) for q in q_norms], (-1, 1))

66 weights_coh[at] = weight_array

67 if sample.has_incoherent and weights.supports_incoherent:

68 weights_incoh = dict()

69 for at in sample.atom_types:

70 weight_array = np.reshape([weights.get_weight_incoh(at, q) for q in q_norms], (-1, 1))

71 weights_incoh[at] = weight_array

73 # weighting of correlation functions

74 if isinstance(sample, StaticSample):

75 data_dict_Sq = _compute_weighting_coherent(sample, 'Sq', weights_coh)

76 data_dict.update(data_dict_Sq)

77 elif isinstance(sample, DynamicSample): 77 ↛ 108line 77 didn't jump to line 108, because the condition on line 77 was never false

78 # coherent

79 Fqt_coh_dict = _compute_weighting_coherent(sample, 'Fqt_coh', weights_coh)

80 data_dict.update(Fqt_coh_dict)

81 Sqw_coh_dict = _compute_weighting_coherent(sample, 'Sqw_coh', weights_coh)

82 data_dict.update(Sqw_coh_dict)

84 # incoherent

85 if sample.has_incoherent and weights.supports_incoherent:

86 Fqt_incoh_dict = _compute_weighting_incoherent(sample, 'Fqt_incoh', weights_incoh)

87 data_dict.update(Fqt_incoh_dict)

88 Sqw_incoh_dict = _compute_weighting_incoherent(sample, 'Sqw_incoh', weights_incoh)

89 data_dict.update(Sqw_incoh_dict)

90 data_dict['Fqt'] = data_dict['Fqt_coh'] + data_dict['Fqt_incoh']

91 data_dict['Sqw'] = data_dict['Sqw_coh'] + data_dict['Sqw_incoh']

92 else:

93 data_dict['Fqt'] = data_dict['Fqt_coh'].copy()

94 data_dict['Sqw'] = data_dict['Sqw_coh'].copy()

96 # currents

97 if sample.has_currents and weights.supports_currents:

98 Clqt_dict = _compute_weighting_coherent(sample, 'Clqt', weights_coh)

99 data_dict.update(Clqt_dict)

100 Clqw_dict = _compute_weighting_coherent(sample, 'Clqw', weights_coh)

101 data_dict.update(Clqw_dict)

103 Ctqt_dict = _compute_weighting_coherent(sample, 'Ctqt', weights_coh)

104 data_dict.update(Ctqt_dict)

105 Ctqw_dict = _compute_weighting_coherent(sample, 'Ctqw', weights_coh)

106 data_dict.update(Ctqw_dict)

108 return Sample(data_dict=data_dict, **meta_data)

111def _compute_weighting_coherent(sample: Sample, name: str, weight_dict: Dict):

112 """

113 Helper function for weighting and summing partial coherent correlation functions.

114 """

115 data_dict = dict()

116 total = np.zeros(sample[name].shape)

117 for s1, s2 in sample.pairs:

118 key_pair = f'{name}_{s1}_{s2}'

119 partial = weight_dict[s1] * weight_dict[s2] * sample[key_pair]

120 data_dict[key_pair] = partial

121 total += partial

122 data_dict[name] = total

123 return data_dict

126def _compute_weighting_incoherent(sample: Sample, name: str, weight_dict: Dict):

127 """

128 Helper function for weighting and summing partial incoherent correlation functions.

129 """

130 data_dict = dict()

131 total = np.zeros(sample[name].shape)

132 for s1 in sample.atom_types:

133 key = f'{name}_{s1}'

134 partial = weight_dict[s1] * sample[key]

135 data_dict[key] = partial

136 total += partial

137 data_dict[name] = total

138 return data_dict