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

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 

7 

8 

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

10 r""" 

11 Weights correlation functions with atomic weighting factors 

12 

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 

16 

17 .. math:: 

18 

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) 

22 

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. 

25 

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). 

29 

30 Parameters 

31 ---------- 

32 sample 

33 input sample to be weighted 

34 weights 

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

36 

37 Returns 

38 ------- 

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

40 """ 

41 

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.') 

46 

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.') 

50 

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) 

56 

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) 

62 

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 

72 

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) 

83 

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() 

95 

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) 

102 

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) 

107 

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

109 

110 

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 

124 

125 

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