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, 

10 weights: Weights, 

11 atom_type_map: Dict[str, str] = None) -> Sample: 

12 r""" 

13 Weights correlation functions with atomic weighting factors 

14 

15 The weighting of a partial dynamic structure factor 

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

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

18 

19 .. math:: 

20 

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

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

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

24 

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

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

27 

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

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

30 returned :attr:`Sample` object (and analogously for current correlation functions). 

31 

32 Parameters 

33 ---------- 

34 sample 

35 input sample to be weighted 

36 weights 

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

38 atom_type_map 

39 map between the atom types in the `Sample` and the ones used in 

40 the `Weights`, e.g., Ba->Ba2+. 

41 

42 Returns 

43 ------- 

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

45 """ 

46 

47 # check input arguments 

48 if sample.has_incoherent and not weights.supports_incoherent: 

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

50 'from the weighted sample.') 

51 

52 if sample.has_currents and not weights.supports_currents: 

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

54 'from the weighted sample.') 

55 

56 # setup new input dicts for new Sample 

57 data_dict = dict() 

58 for key in sample.dimensions: 

59 data_dict[key] = sample[key] 

60 meta_data = deepcopy(sample.meta_data) 

61 

62 # Map the atom types in the sample to the types in the weights object. 

63 # Useful, for instance, when using weights for charged atomic species. 

64 atom_types = [(at, at) for at in sample.atom_types] 

65 if atom_type_map is not None: 

66 # Fallback to use the atom type from species (`at`) if it's not mapped. 

67 atom_types = [(at, atom_type_map.get(at, at)) for at in atom_type_map] 

68 

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

70 if 'q_norms' in sample.dimensions: 

71 q_norms = sample.q_norms 

72 else: 

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

74 

75 weights_coh = dict() 

76 for at, weight_at in atom_types: 

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

78 weights_coh[at] = weight_array 

79 if sample.has_incoherent and weights.supports_incoherent: 

80 weights_incoh = dict() 

81 for at, weight_at in atom_types: 

82 weight_array = np.reshape([ 

83 weights.get_weight_incoh(weight_at, q) for q in q_norms 

84 ], (-1, 1)) 

85 weights_incoh[at] = weight_array 

86 

87 # weighting of correlation functions 

88 if isinstance(sample, StaticSample): 

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

90 data_dict.update(data_dict_Sq) 

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

92 # coherent 

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

94 data_dict.update(Fqt_coh_dict) 

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

96 data_dict.update(Sqw_coh_dict) 

97 

98 # incoherent 

99 if sample.has_incoherent and weights.supports_incoherent: 

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

101 data_dict.update(Fqt_incoh_dict) 

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

103 data_dict.update(Sqw_incoh_dict) 

104 

105 # currents 

106 if sample.has_currents and weights.supports_currents: 

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

108 data_dict.update(Clqt_dict) 

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

110 data_dict.update(Clqw_dict) 

111 

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

113 data_dict.update(Ctqt_dict) 

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

115 data_dict.update(Ctqw_dict) 

116 

117 return sample.__class__(data_dict, **meta_data) 

118 

119 

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

121 """ 

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

123 """ 

124 data_dict = dict() 

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

126 for s1, s2 in sample.pairs: 

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

128 partial = np.real(np.conjugate(weight_dict[s1]) * weight_dict[s2]) * sample[key_pair] 

129 data_dict[key_pair] = partial 

130 total += partial 

131 data_dict[name] = total 

132 return data_dict 

133 

134 

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

136 """ 

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

138 """ 

139 data_dict = dict() 

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

141 for s1 in sample.atom_types: 

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

143 partial = np.real(np.conjugate(weight_dict[s1]) * weight_dict[s1]) * sample[key] 

144 data_dict[key] = partial 

145 total += partial 

146 data_dict[name] = total 

147 return data_dict