Coverage for dynasor/post_processing/weights.py: 100%

1from typing import Any, Optional 

2from pandas import DataFrame 

3 

4 

5class Weights: 

6 """ 

7 Class holding weights and support functions for weighting of samples. 

8 

9 Parameters 

10 ---------- 

11 weights_coh 

12 A dict with keys and values representing the atom types and their corresponding 

13 coherent scattering length, ``{'A': b_A }``. 

14 weights_incoh 

15 A dict with keys and values representing the atom types and their corresponding 

16 incoherent scattering length, ``{'A': b_A }``. 

17 supports_currents 

18 Whether or not the coherent weights should be applied to current-correlation functions. 

19 """ 

20 

21 def __init__( 

22 self, 

23 weights_coh: dict[str, Any], 

24 weights_incoh: Optional[dict[str, Any]] = None, 

25 supports_currents: Optional[bool] = True, 

26 ): 

27 if weights_incoh is not None: 

28 if set(weights_coh.keys()) != set(weights_incoh.keys()): 

29 raise ValueError( 

30 'Incoherent weights keys do not match coherent weights keys: ' 

31 f'{sorted(weights_incoh.keys())} vs {sorted(weights_coh.keys())}' 

32 ) 

33 self._weights_coh = weights_coh 

34 self._weights_incoh = weights_incoh 

35 self._supports_currents = supports_currents 

36 

37 def get_weight_coh(self, atom_type: str, q_norm: float = None) -> float: 

38 """Get the coherent weight for a given atom type and q-vector norm.""" 

39 if atom_type not in self._weights_coh.keys(): 

40 raise ValueError(f'Coherent weights for {atom_type} have not been specified') 

41 return self._weights_coh[atom_type] 

42 

43 def get_weight_incoh(self, atom_type: str, q_norm: float = None) -> float: 

44 """Get the incoherent weight for a given atom type and q-vector norm.""" 

45 if self._weights_incoh is None: 

46 return None 

47 if atom_type not in self._weights_incoh.keys(): 

48 raise ValueError(f'Incoherent weights for {atom_type} have not been specified') 

49 return self._weights_incoh[atom_type] 

50 

51 @property 

52 def parameters(self) -> DataFrame: 

53 """Parameters used for weighting. 

54 """ 

55 atom_types = set(self._weights_coh.keys()) 

56 if self._weights_incoh is not None: 

57 atom_types |= set(self._weights_incoh.keys()) 

58 data = [] 

59 for s in atom_types: 

60 record = dict( 

61 atom_type=s, 

62 coherent=self._weights_coh.get(s, None), 

63 ) 

64 data.append(record) 

65 if self._weights_incoh is not None: 

66 record.update(dict(incoherent=self._weights_incoh.get(s, None))) 

67 return DataFrame(data).sort_values('atom_type', ignore_index=True) 

68 

69 @property 

70 def supports_currents(self) -> bool: 

71 """ 

72 Whether or not this :class:`Weights` object supports weighting of current correlations. 

73 """ 

74 return self._supports_currents 

75 

76 @property 

77 def supports_incoherent(self) -> bool: 

78 """ 

79 Whether or not this :class:`Weights` object supports weighting of incoherent 

80 correlation functions. 

81 """ 

82 return self._weights_incoh is not None 

83 

84 def __str__(self): 

85 s = ['weights coherent:'] 

86 for key, val in self._weights_coh.items(): 

87 s.append(f' {key}: {val}') 

88 

89 # Return early if incoherent weights 

90 # are None 

91 if self._weights_incoh is None: 

92 return '\n'.join(s) 

93 

94 s.append('weights incoherent:') 

95 for key, val in self._weights_incoh.items(): 

96 s.append(f' {key}: {val}') 

97 return '\n'.join(s)