Coverage for local_installation/dynasor/sample.py: 92%

1import numpy as np 

2import pandas as pd 

3from typing import Dict, Any 

4 

5 

6class Sample: 

7 """ 

8 Class for holding correlation functions and some additional meta data. 

9 Sample objects can be written to and read from file. 

10 

11 Parameters 

12 ---------- 

13 data_dict 

14 Dictionary with correlation functions. 

15 meta_data 

16 Dictionary with meta data, for example atom-types, simulation cell, number of atoms, 

17 time stamps, user names, etc. 

18 """ 

19 

20 def __init__(self, data_dict: Dict[str, Any], **meta_data: Dict[str, Any]): 

21 

22 # set data dict as attributes 

23 self._data_keys = list(data_dict) 

24 for key in data_dict: 

25 setattr(self, key, data_dict[key]) 

26 

27 # set system parameters 

28 self.meta_data = meta_data 

29 self._atom_types = meta_data['atom_types'] 

30 self._pairs = meta_data['pairs'] 

31 self._particle_counts = meta_data['particle_counts'] 

32 self._cell = meta_data['cell'] 

33 

34 def __getitem__(self, key): 

35 """ Makes it possible to get the attributes using Sample['key'] """ 

36 try: 

37 return getattr(self, key) 

38 except AttributeError: 

39 raise KeyError(key) 

40 

41 def write_to_npz(self, fname: str): 

42 """ Write object to file in numpy's npz format. 

43 

44 Parameters 

45 ---------- 

46 fname 

47 Name of the file in which to store the Sample object. 

48 """ 

49 data_to_save = dict(name=self.__class__.__name__) 

50 data_to_save['meta_data'] = self.meta_data 

51 data_dict = dict() 

52 for key in self._data_keys: 

53 data_dict[key] = getattr(self, key) 

54 data_to_save['data_dict'] = data_dict 

55 np.savez_compressed(fname, **data_to_save) 

56 

57 @property 

58 def available_correlation_functions(self): 

59 """ All the available correlation functions in sample. """ 

60 keys_to_skip = set(['q_points', 'q_norms', 'time', 'omega']) 

61 return sorted(list(set(self._data_keys) - keys_to_skip)) 

62 

63 @property 

64 def dimensions(self): 

65 r"""The dimensions for the samples, e.g., for 

66 :math:`S(q, \omega)` the dimensions of the 

67 :math:`q` and :math:`\omega` axes. 

68 """ 

69 keys_to_skip = set(self.available_correlation_functions) 

70 return sorted(list(set(self._data_keys) - keys_to_skip)) 

71 

72 @property 

73 def atom_types(self): 

74 if self._atom_types is None: 74 ↛ 75line 74 didn't jump to line 75, because the condition on line 74 was never true

75 return None 

76 return self._atom_types.copy() 

77 

78 @property 

79 def particle_counts(self): 

80 if self._particle_counts is None: 80 ↛ 81line 80 didn't jump to line 81, because the condition on line 80 was never true

81 return None 

82 return self._particle_counts.copy() 

83 

84 @property 

85 def pairs(self): 

86 if self._pairs is None: 86 ↛ 87line 86 didn't jump to line 87, because the condition on line 86 was never true

87 return None 

88 return self._pairs.copy() 

89 

90 @property 

91 def cell(self): 

92 if self._cell is None: 92 ↛ 93line 92 didn't jump to line 93, because the condition on line 92 was never true

93 return None 

94 return self._cell.copy() 

95 

96 def __repr__(self): 

97 return str(self) 

98 

99 def __str__(self): 

100 s_contents = [self.__class__.__name__] 

101 s_contents.append(f'Atom types: {self.atom_types}') 

102 s_contents.append(f'Pairs: {self.pairs}') 

103 s_contents.append(f'Particle counts: {self.particle_counts}') 

104 s_contents.append('Simulations cell:') 

105 s_contents.append(f'{self.cell}') 

106 for key in self._data_keys: 

107 s_i = f'{key:15} with shape: {np.shape(getattr(self,key))}' 

108 s_contents.append(s_i) 

109 s = '\n'.join(s_contents) 

110 return s 

111 

112 def _repr_html_(self) -> str: 

113 s = [f'<h3>{self.__class__.__name__}</h3>'] 

114 s += ['<table border="1" class="dataframe">'] 

115 s += ['<thead><tr><th style="text-align: left">Field</th>' 

116 '<th>Content/Size</th></tr></thead>'] 

117 s += ['<tbody>'] 

118 s += ['<tr><td style="text-align: left">Atom types</td>' 

119 f'<td>{self.atom_types}</td></tr>'] 

120 s += ['<tr><td style="text-align: left">Pairs</td>' 

121 f'<td>{self.pairs}</td></tr>'] 

122 s += ['<tr><td style="text-align: left">Particle counts</td>' 

123 f'<td>{self.particle_counts}</td></tr>'] 

124 s += ['<tr><td style="text-align: left">Simulations cell</td>' 

125 f'<td>{self.cell}</td></tr>'] 

126 for key in self._data_keys: 

127 s += [f'<tr><td style="text-align: left">{key}</td>' 

128 f'<td>{np.shape(getattr(self,key))}</td></tr>'] 

129 s += ['</tbody>'] 

130 s += ['</table>'] 

131 return '\n'.join(s) 

132 

133 @property 

134 def has_incoherent(self): 

135 """ Whether this sample contains the incoherent correlation functions or not. """ 

136 return False 

137 

138 @property 

139 def has_currents(self): 

140 """ Whether this sample contains the current correlation functions or not. """ 

141 return False 

142 

143 

144class StaticSample(Sample): 

145 

146 def to_dataframe(self): 

147 """ Returns correlation functions as pandas dataframe """ 

148 df = pd.DataFrame() 

149 for dim in self.dimensions: 

150 df[dim] = self[dim].tolist() # to list to make q-points (N, 3) work in dataframe 

151 for key in self.available_correlation_functions: 

152 df[key] = self[key].reshape(-1, ) 

153 return df 

154 

155 

156class DynamicSample(Sample): 

157 

158 @property 

159 def has_incoherent(self): 

160 return 'Fqt_incoh' in self.available_correlation_functions 

161 

162 @property 

163 def has_currents(self): 

164 pair_string = '_'.join(self.pairs[0]) 

165 return f'Clqt_{pair_string}' in self.available_correlation_functions 

166 

167 def to_dataframe(self, q_index: int): 

168 """ Returns correlation functions as pandas dataframe for the given q-index. 

169 

170 Parameters 

171 ---------- 

172 q_index 

173 index of q-point to return 

174 """ 

175 df = pd.DataFrame() 

176 for dim in self.dimensions: 

177 if dim in ['q_points', 'q_norms']: 

178 continue 

179 df[dim] = self[dim] 

180 for key in self.available_correlation_functions: 

181 df[key] = self[key][q_index] 

182 return df 

183 

184 

185def read_sample_from_npz(fname: str) -> Sample: 

186 """ Read :class:`Sample <dynasor.sample.Sample>` from file. 

187 

188 Parameters 

189 ---------- 

190 fname 

191 Path to the file (numpy npz format) from which to read 

192 the :class:`Sample <dynasor.sample.Sample>` object. 

193 """ 

194 data_read = np.load(fname, allow_pickle=True) 

195 data_dict = data_read['data_dict'].item() 

196 meta_data = data_read['meta_data'].item() 

197 if data_read['name'] == 'StaticSample': 197 ↛ 198line 197 didn't jump to line 198, because the condition on line 197 was never true

198 return StaticSample(data_dict, **meta_data) 

199 elif data_read['name'] == 'DynamicSample': 

200 return DynamicSample(data_dict, **meta_data) 

201 else: 

202 return Sample(data_dict, **meta_data)