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

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 

8 

9 

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 

15 

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 

19 

20 .. math:: 

21 

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) 

25 

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. 

28 

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 :attr:`Sample` object (and analogously for current correlation functions). 

32 

33 Parameters 

34 ---------- 

35 sample 

36 Input sample to be weighted. 

37 weights 

38 Object containing the weights :math:`f_\mathrm{X}(\boldsymbol{q})`. 

39 atom_type_map 

40 Map between the atom types in the :class:`Sample` and the ones used in 

41 the :class:`Weights` object, e.g., Ba → Ba(2+). 

42 

43 Returns 

44 ------- 

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

46 """ 

47 

48 # check input arguments 

49 if not isinstance(sample, (StaticSample, DynamicSample)): 

50 raise TypeError('sample must be a StaticSample or DynamicSample.') 

51 if sample.has_incoherent and not weights.supports_incoherent: 

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

53 'from the weighted sample.') 

54 

55 if sample.has_currents and not weights.supports_currents: 

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

57 'from the weighted sample.') 

58 

59 # setup new input dicts for new Sample 

60 data_dict = dict() 

61 for key in sample.dimensions: 

62 data_dict[key] = sample[key] 

63 

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

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

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

67 if atom_type_map is not None: 

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

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

70 

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

72 if 'q_norms' in sample.dimensions: 

73 q_norms = sample.q_norms 

74 else: 

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

76 

77 weights_coh = dict() 

78 for at, weight_at in atom_types: 

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

80 weights_coh[at] = weight_array 

81 if sample.has_incoherent and weights.supports_incoherent: 

82 weights_incoh = dict() 

83 for at, weight_at in atom_types: 

84 weight_array = np.reshape([ 

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

86 ], (-1, 1)) 

87 weights_incoh[at] = weight_array 

88 

89 # weighting of correlation functions 

90 if isinstance(sample, StaticSample): 

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

92 data_dict.update(data_dict_Sq) 

93 elif isinstance(sample, DynamicSample): 93 ↛ 119line 93 didn't jump to line 119 because the condition on line 93 was always true

94 # coherent 

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

96 data_dict.update(Fqt_coh_dict) 

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

98 data_dict.update(Sqw_coh_dict) 

99 

100 # incoherent 

101 if sample.has_incoherent and weights.supports_incoherent: 

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

103 data_dict.update(Fqt_incoh_dict) 

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

105 data_dict.update(Sqw_incoh_dict) 

106 

107 # currents 

108 if sample.has_currents and weights.supports_currents: 

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

110 data_dict.update(Clqt_dict) 

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

112 data_dict.update(Clqw_dict) 

113 

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

115 data_dict.update(Ctqt_dict) 

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

117 data_dict.update(Ctqw_dict) 

118 

119 new_sample = sample.__class__( 

120 data_dict, 

121 simulation_data=deepcopy(sample.simulation_data), 

122 history=deepcopy(sample.history)) 

123 new_sample._append_history( 

124 'get_weighted_sample', 

125 dict( 

126 atom_type_map=atom_type_map, 

127 weights_class=weights.__class__.__name__, 

128 weights_parameters=weights.parameters.to_dict(), 

129 )) 

130 

131 return new_sample 

132 

133 

134def _compute_weighting_coherent( 

135 sample: Sample, 

136 name: str, 

137 weight_dict: dict, 

138) -> dict[str, NDArray[float]]: 

139 """ 

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

141 """ 

142 data_dict = dict() 

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

144 for s1, s2 in sample.pairs: 

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

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

147 data_dict[key_pair] = partial 

148 total += partial 

149 data_dict[name] = total 

150 return data_dict 

151 

152 

153def _compute_weighting_incoherent( 

154 sample: Sample, 

155 name: str, 

156 weight_dict: dict, 

157) -> dict[str, NDArray[float]]: 

158 """ 

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

160 """ 

161 data_dict = dict() 

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

163 for s1 in sample.atom_types: 

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

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

166 data_dict[key] = partial 

167 total += partial 

168 data_dict[name] = total 

169 return data_dict