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

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 self._weights_coh = weights_coh 

28 self._weights_incoh = weights_incoh 

29 self._supports_currents = supports_currents 

30 

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

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

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

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

35 return self._weights_coh[atom_type] 

36 

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

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

39 if self._weights_incoh is None: 

40 return None 

41 if atom_type not in self._weights_incoh.keys(): 41 ↛ 42line 41 didn't jump to line 42 because the condition on line 41 was never true

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

43 return self._weights_incoh[atom_type] 

44 

45 @property 

46 def parameters(self) -> DataFrame: 

47 """Parameters used for weighting. 

48 """ 

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

50 if self._weights_incoh: 

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

52 data = [] 

53 for s in atom_types: 

54 record = dict( 

55 atom_type=s, 

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

57 ) 

58 data.append(record) 

59 if self._weights_incoh: 

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

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

62 

63 @property 

64 def supports_currents(self) -> bool: 

65 """ 

66 Wether or not this :class:`Weights` object supports weighting of current correlations. 

67 """ 

68 return self._supports_currents 

69 

70 @property 

71 def supports_incoherent(self) -> bool: 

72 """ 

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

74 correlation functions. 

75 """ 

76 return self._weights_incoh is not None 

77 

78 def __str__(self): 

79 s = ['weights coherent:'] 

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

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

82 

83 # Return early if incoherent weights 

84 # are None 

85 if self._weights_incoh is None: 

86 return '\n'.join(s) 

87 

88 s.append('weights incoherent:') 

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

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

91 return '\n'.join(s)