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