Coverage for dynasor/sample.py: 100%
173 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-10 08:27 +0000
1import getpass
2import socket
3from copy import deepcopy
4from datetime import datetime
5from typing import Any, Optional
6import numpy as np
7from dynasor.logging_tools import logger
8from numpy.typing import NDArray
9from pandas import DataFrame
12class Sample:
13 """
14 Class for holding correlation functions and additional metadata.
15 Objects of this class are most commonly generated by calling functions such
16 as :func:`compute_static_structure_factors` or
17 :func:`compute_dynamic_structure_factors`.
18 They can then be written to and subsequently read from file.
20 You can see which correlation functions are available via the
21 :attr:`available_correlation_functions` property.
22 You can then access the correlation functions either by key or as property.
23 For example, you could access the static structure factor in the following ways::
25 sample.Sq # as property
26 sample['Sq'] # via key
28 The correlation functions are provided as numpy arrays.
30 There are several additional fields, the availability of which depends on the
31 type of correlation function that was sampled. Static samples, for example, do
32 not contain the `time` and `omega` fields. `q_norms` is only available if spherical
33 averaging was carried out, typically via :func:`get_spherically_averaged_sample_smearing`.
35 * `q_points`: list of q-point coordinates
36 * `q_norms`: norms of the momentum vector
37 * `time`: time
38 * `omega`: frequency
40 You can also see which fields are available by "printing" the :class:`Sample` object.
42 Parameters
43 ----------
44 data_dict
45 Dictionary with correlation functions.
46 simulation_data
47 Dictionary with simulation data. The following fields are strongly encouraged
48 (but not enforced): `atom_types`, `cell`, `particle_counts`.
49 history
50 Previous history of operations on :class:`Sample` object.
51 """
53 def __init__(
54 self,
55 data_dict: dict[str, Any],
56 simulation_data: dict[str, Any],
57 history: Optional[list[dict[str, Any]]] = None,
58 ):
59 reserved_keys = set(dir(type(self))) | {'_data_keys', '_metadata'}
60 colliding_keys = sorted(reserved_keys.intersection(data_dict))
61 if colliding_keys:
62 raise ValueError(f'data_dict contains reserved Sample field names: {colliding_keys}')
64 # set data dict as attributes
65 self._data_keys = list(data_dict)
66 for key in data_dict:
67 setattr(self, key, data_dict[key])
69 # set metadata
70 # (using deepcopy here to avoid accidental transfer by reference)
71 self._metadata = dict()
72 self._metadata['simulation_data'] = deepcopy(simulation_data)
73 if history is not None:
74 logger.debug('Copying history')
75 self._metadata['history'] = deepcopy(history)
76 else:
77 self._metadata['history'] = []
79 def _append_history(
80 self,
81 calling_function: str,
82 caller_metadata: Optional[dict[str, Any]] = None,
83 ):
84 """Add record to history.
86 Parameters
87 ----------
88 calling_function
89 Name of calling function.
90 caller_metadata
91 Metadata associated with the calling function.
92 """
93 from dynasor import __version__ as dynasor_version
95 new_record = dict(func=calling_function)
96 if caller_metadata is not None:
97 new_record.update(caller_metadata.copy())
98 new_record.update(dict(
99 date_time=datetime.now().strftime('%Y-%m-%dT%H:%M:%S'),
100 username=getpass.getuser(),
101 hostname=socket.gethostname(),
102 dynasor_version=dynasor_version,
103 ))
104 self._metadata['history'].append(new_record)
106 def __getitem__(self, key):
107 """ Makes it possible to get the attributes using Sample['key'] """
108 try:
109 return getattr(self, key)
110 except AttributeError:
111 raise KeyError(key)
113 def write_to_npz(self, fname: str):
114 """ Write object to file in numpy npz format.
116 Parameters
117 ----------
118 fname
119 Name of the file in which to store the Sample object.
120 """
121 data_to_save = dict(name=self.__class__.__name__)
122 data_to_save['metadata'] = self._metadata
123 data_dict = dict()
124 for key in self._data_keys:
125 data_dict[key] = getattr(self, key)
126 data_to_save['data_dict'] = data_dict
127 np.savez_compressed(fname, **data_to_save)
129 @property
130 def available_correlation_functions(self) -> list[str]:
131 """ All the available correlation functions in sample. """
132 keys_to_skip = set(['q_points', 'q_norms', 'time', 'omega'])
133 return sorted(list(set(self._data_keys) - keys_to_skip))
135 @property
136 def dimensions(self) -> list[str]:
137 r"""The dimensions for the samples, e.g., for :math:`S(q, \omega)`
138 the dimensions would be the :math:`q` and :math:`\omega` axes.
139 """
140 keys_to_skip = set(self.available_correlation_functions)
141 return sorted(list(set(self._data_keys) - keys_to_skip))
143 @property
144 def metadata(self) -> dict[str, Any]:
145 """ Metadata. """
146 return deepcopy(self._metadata)
148 @property
149 def simulation_data(self) -> dict[str, Any]:
150 """ Simulation data. """
151 return deepcopy(self._metadata['simulation_data'])
153 @property
154 def history(self) -> list[dict[str, Any]]:
155 """ List of operations applied to this :class:`Sample` object. """
156 return deepcopy(self._metadata['history'])
158 @property
159 def atom_types(self) -> list[str]:
160 """ Simulation data: Atom types. """
161 return self.simulation_data['atom_types'].copy() \
162 if 'atom_types' in self.simulation_data else None
164 @property
165 def particle_counts(self) -> dict[str, int]:
166 """ Simulation data: Number of particles per type. """
167 return self.simulation_data['particle_counts'].copy() \
168 if 'particle_counts' in self.simulation_data else None
170 @property
171 def pairs(self) -> list[tuple[str, str]]:
172 """ Pairs of types for which correlation functions are available. """
173 return self.simulation_data['pairs'].copy() \
174 if 'pairs' in self.simulation_data else None
176 @property
177 def cell(self) -> NDArray[float]:
178 """ Simulation data: Cell metric. """
179 return self.simulation_data['cell'].copy() \
180 if 'cell' in self.simulation_data else None
182 @property
183 def has_incoherent(self) -> bool:
184 """ Whether this sample contains the incoherent correlation functions or not. """
185 return False
187 @property
188 def has_currents(self) -> bool:
189 """ Whether this sample contains the current correlation functions or not. """
190 return False
192 def __repr__(self) -> str:
193 return str(self)
195 _special_fields = dict(
196 atom_types='Atom types',
197 cell='Cell',
198 number_of_frames='Number of frames',
199 particle_counts='Particle counts',
200 maximum_time_lag='Maximum time lag',
201 angular_frequency_resolution='Angular frequency resolution',
202 time_between_frames='Time between frames',
203 )
205 def __str__(self) -> str:
206 s_contents = [self.__class__.__name__]
207 for key, value in sorted(self.simulation_data.items()):
208 s_contents.append(f'{self._special_fields.get(key, key)}: {value}')
209 for key in self.dimensions:
210 s_i = f'{key:15} with shape: {np.shape(getattr(self, key))}'
211 s_contents.append(s_i)
212 for key in self.available_correlation_functions:
213 s_i = f'{key:15} with shape: {np.shape(getattr(self, key))}'
214 s_contents.append(s_i)
215 s = '\n'.join(s_contents)
216 return s
218 def _repr_html_(self) -> str:
219 s = [f'<h3>{self.__class__.__name__}</h3>']
221 s += ['<h4>Simulation</h4>']
222 s += ['<table border="1" class="dataframe">']
223 s += ['<thead><tr>'
224 '<th style="text-align: left">Name</th>'
225 '<th>Content</th>'
226 '</tr></thead>']
227 s += ['<tbody>']
228 for key, value in sorted(self.simulation_data.items()):
229 if key not in self._special_fields:
230 continue
231 s += [f'<tr><td style="text-align: left;">{self._special_fields[key]}</td>'
232 f'<td>{value}</td></tr>']
233 s += ['</tbody>']
234 s += ['</table>']
236 s += ['<h4>Dimensions</h4>']
237 s += ['<table border="1" class="dataframe">']
238 s += ['<thead><tr>'
239 '<th style="text-align: left">Field</th>'
240 '<th>Size</th>'
241 '</tr></thead>']
242 s += ['<tbody>']
243 for key in self.dimensions:
244 s += [f'<tr><td style="text-align: left">{key}</td>'
245 f'<td>{np.shape(getattr(self, key))}</td></tr>']
246 s += ['</tbody>']
247 s += ['</table>']
249 s += ['<h4>History</h4>']
250 s += ['<table border="1" class="dataframe">']
251 s += ['<thead><tr>'
252 '<th style="text-align: left">Function</th>'
253 '<th style="text-align: left">Field</th>'
254 '<th style="text-align: left">Content</th>'
255 '</tr></thead>']
256 s += ['<tbody>']
257 for entry in self.history:
258 title = entry.get('func', '')
259 for key, value in entry.items():
260 if key == 'func':
261 continue
262 s += [f'<tr><td style="text-align: left">{title}</td>'
263 f'<td style="text-align: left">{key}</td>'
264 f'<td>{value}</td></tr>']
265 title = ''
266 s += ['</tbody>']
267 s += ['</table>']
268 return '\n'.join(s)
271class StaticSample(Sample):
272 """
273 Class for holding static correlation functions and additional metadata.
274 Objects of this class are most commonly generated by calling
275 :func:`compute_static_structure_factors <dynasor.compute_static_structure_factors>`.
276 They can then be written to and subsequently read from file.
278 You can see which correlation functions are available via the
279 :attr:`available_correlation_functions` property.
280 You can then access the correlation functions either by key or as property.
281 For example, you could access the static structure factor :math:`S(q)`
282 in the following ways::
284 sample.Sq # as property
285 sample['Sq'] # via key
287 The correlation functions are provided as numpy arrays.
289 There are several additional fields:
291 * `q_points`: list of q-point coordinates
292 * `q_norms`: norms of the momentum vector (available when the :class:`StaticSample`
293 object was generated via :func:`get_spherically_averaged_sample_smearing
294 <dynasor.post_processing.get_spherically_averaged_sample_smearing>` or similar
295 functions)
297 You can also see which fields are available by "printing" the :class:`StaticSample` object.
299 Parameters
300 ----------
301 data_dict
302 Dictionary with correlation functions.
303 simulation_data
304 Dictionary with simulation data. The following fields are strongly encouraged
305 (but not enforced): `atom_types`, `cell`, `particle_counts`.
306 history
307 Previous history of operations on :class:`Sample` object.
308 """
310 def to_dataframe(self) -> DataFrame:
311 """ Returns correlation functions as pandas dataframe """
312 df = DataFrame()
313 for dim in self.dimensions:
314 df[dim] = self[dim].tolist() # to list to make q-points (N, 3) work in dataframe
315 for key in self.available_correlation_functions:
316 df[key] = self[key].reshape(-1, )
317 return df
320class DynamicSample(Sample):
321 r"""
322 Class for holding dynamic correlation functions and additional metadata.
323 Objects of this class are most commonly generated by calling
324 :func:`compute_dynamic_structure_factors <dynasor.compute_dynamic_structure_factors>`.
325 They can then be written to and subsequently read from file.
327 You can see which correlation functions are available via the
328 :attr:`available_correlation_functions` property.
329 You can then access the correlation functions either by key or as property.
330 For example, you could access the dynamic structure factor :math:`S(q,\omega)`
331 in the following ways::
333 sample.Sqw # as property
334 sample['Sqw'] # via key
336 The correlation functions are provided as numpy arrays.
338 There are several additional fields, the availability of which depends on the
339 type of correlation function that was sampled.
341 * `q_points`: list of q-point coordinates
342 * `q_norms`: norms of the momentum vector (available when the :class:`DynamicSample`
343 object was generated, e.g., via :func:`get_spherically_averaged_sample_smearing
344 <dynasor.post_processing.get_spherically_averaged_sample_smearing>` or similar
345 functions)
346 * `time`: time
347 * `omega`: frequency
349 You can also see which fields are available by "printing" the :class:`DynamicSample` object.
351 Parameters
352 ----------
353 data_dict
354 Dictionary with correlation functions.
355 simulation_data
356 Dictionary with simulation data. The following fields are strongly encouraged
357 (but not enforced): `atom_types`, `cell`, `particle_counts`.
358 history
359 Previous history of operations on :class:`Sample` object.
360 """
362 @property
363 def has_incoherent(self) -> bool:
364 return 'Fqt_incoh' in self.available_correlation_functions
366 @property
367 def has_currents(self) -> bool:
368 pair_string = '_'.join(self.pairs[0])
369 return f'Clqt_{pair_string}' in self.available_correlation_functions
371 def to_dataframe(self, q_index: int) -> DataFrame:
372 """ Returns correlation functions as pandas dataframe for the given q-index.
374 Parameters
375 ----------
376 q_index
377 Index of q-point to return.
378 """
379 df = DataFrame()
380 for dim in self.dimensions:
381 if dim in ['q_points', 'q_norms']:
382 continue
383 df[dim] = self[dim]
384 for key in self.available_correlation_functions:
385 df[key] = self[key][q_index]
386 return df
389def read_sample_from_npz(fname: str) -> Sample:
390 """ Read :class:`Sample <dynasor.sample.Sample>` from file.
392 Parameters
393 ----------
394 fname
395 Path to the file (numpy npz format) from which to read
396 the :class:`Sample <dynasor.sample.Sample>` object.
397 """
398 with np.load(fname, allow_pickle=True) as data_read:
399 try:
400 metadata = data_read['metadata'].item()
401 except KeyError:
402 # fallback for versions<=2.2
403 metadata = data_read['meta_data'].item()
404 data_dict = data_read['data_dict'].item()
405 sample_name = data_read['name'].item()
407 if 'simulation_data' in metadata:
408 logger.debug(f'Reading Sample object from {fname} assuming version >=2.3')
409 simulation_data = metadata['simulation_data']
410 else:
411 logger.debug(f'Reading Sample object from {fname} assuming version <=2.2')
412 simulation_data = {}
413 for key in [
414 'atom_types', 'pairs', 'particle_counts', 'cell',
415 'time_between_frames', 'maximum_time_lag', 'angular_frequency_resolution',
416 'maximum_angular_frequency', 'number_of_frames',
417 ]:
418 if key in metadata:
419 simulation_data[key] = metadata[key]
421 history = metadata['history'] if 'history' in metadata else None
423 if sample_name == 'StaticSample':
424 return StaticSample(data_dict, simulation_data, history=history)
425 elif sample_name == 'DynamicSample':
426 return DynamicSample(data_dict, simulation_data, history=history)
427 else:
428 return Sample(data_dict, simulation_data, history=history)