Coverage for local_installation/dynasor/sample.py: 92%
120 statements
« prev ^ index » next coverage.py v7.3.2, created at 2023-11-30 21:04 +0000
« prev ^ index » next coverage.py v7.3.2, created at 2023-11-30 21:04 +0000
1import numpy as np
2import pandas as pd
3from typing import Dict, Any
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.
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 """
20 def __init__(self, data_dict: Dict[str, Any], **meta_data: Dict[str, Any]):
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])
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']
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)
41 def write_to_npz(self, fname: str):
42 """ Write object to file in numpy's npz format.
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)
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))
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))
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()
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()
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()
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()
96 def __repr__(self):
97 return str(self)
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
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)
133 @property
134 def has_incoherent(self):
135 """ Whether this sample contains the incoherent correlation functions or not. """
136 return False
138 @property
139 def has_currents(self):
140 """ Whether this sample contains the current correlation functions or not. """
141 return False
144class StaticSample(Sample):
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
156class DynamicSample(Sample):
158 @property
159 def has_incoherent(self):
160 return 'Fqt_incoh' in self.available_correlation_functions
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
167 def to_dataframe(self, q_index: int):
168 """ Returns correlation functions as pandas dataframe for the given q-index.
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
185def read_sample_from_npz(fname: str) -> Sample:
186 """ Read :class:`Sample <dynasor.sample.Sample>` from file.
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)