import numpy as np
import pandas as pd
from typing import Dict, Any
class Sample:
"""
Class for holding correlation functions and some additional meta data.
Sample objects can be written to and read from file.
Parameters
----------
data_dict
Dictionary with correlation functions.
meta_data
Dictionary with meta data, for example atom-types, simulation cell, number of atoms,
time stamps, user names, etc.
"""
def __init__(self, data_dict: Dict[str, Any], **meta_data: Dict[str, Any]):
# set data dict as attributes
self._data_keys = list(data_dict)
for key in data_dict:
setattr(self, key, data_dict[key])
# set system parameters
self.meta_data = meta_data
self._atom_types = meta_data['atom_types']
self._pairs = meta_data['pairs']
self._particle_counts = meta_data['particle_counts']
self._cell = meta_data['cell']
def __getitem__(self, key):
""" Makes it possible to get the attributes using Sample['key'] """
try:
return getattr(self, key)
except AttributeError:
raise KeyError(key)
def write_to_npz(self, fname: str):
""" Write object to file in numpy's npz format.
Parameters
----------
fname
Name of the file in which to store the Sample object.
"""
data_to_save = dict(name=self.__class__.__name__)
data_to_save['meta_data'] = self.meta_data
data_dict = dict()
for key in self._data_keys:
data_dict[key] = getattr(self, key)
data_to_save['data_dict'] = data_dict
np.savez_compressed(fname, **data_to_save)
@property
def available_correlation_functions(self):
""" All the available correlation functions in sample. """
keys_to_skip = set(['q_points', 'q_norms', 'time', 'omega'])
return sorted(list(set(self._data_keys) - keys_to_skip))
@property
def dimensions(self):
r"""The dimensions for the samples, e.g., for :math:`S(q, \omega)`
the dimensions would be the :math:`q` and :math:`\omega` axes.
"""
keys_to_skip = set(self.available_correlation_functions)
return sorted(list(set(self._data_keys) - keys_to_skip))
@property
def atom_types(self):
if self._atom_types is None:
return None
return self._atom_types.copy()
@property
def particle_counts(self):
if self._particle_counts is None:
return None
return self._particle_counts.copy()
@property
def pairs(self):
if self._pairs is None:
return None
return self._pairs.copy()
@property
def cell(self):
if self._cell is None:
return None
return self._cell.copy()
def __repr__(self):
return str(self)
def __str__(self):
s_contents = [self.__class__.__name__]
s_contents.append(f'Atom types: {self.atom_types}')
s_contents.append(f'Pairs: {self.pairs}')
s_contents.append(f'Particle counts: {self.particle_counts}')
s_contents.append('Simulations cell:')
s_contents.append(f'{self.cell}')
for key in self.dimensions:
s_i = f'{key:15} with shape: {np.shape(getattr(self,key))}'
s_contents.append(s_i)
for key in self.available_correlation_functions:
s_i = f'{key:15} with shape: {np.shape(getattr(self,key))}'
s_contents.append(s_i)
s = '\n'.join(s_contents)
return s
def _repr_html_(self) -> str:
s = [f'<h3>{self.__class__.__name__}</h3>']
s += ['<table border="1" class="dataframe">']
s += ['<thead><tr><th style="text-align: left">Field</th>'
'<th>Content/Size</th></tr></thead>']
s += ['<tbody>']
s += ['<tr><td style="text-align: left">Atom types</td>'
f'<td>{self.atom_types}</td></tr>']
s += ['<tr><td style="text-align: left">Pairs</td>'
f'<td>{self.pairs}</td></tr>']
s += ['<tr><td style="text-align: left">Particle counts</td>'
f'<td>{self.particle_counts}</td></tr>']
s += ['<tr><td style="text-align: left">Simulations cell</td>'
f'<td>{self.cell}</td></tr>']
for key in self._data_keys:
s += [f'<tr><td style="text-align: left">{key}</td>'
f'<td>{np.shape(getattr(self,key))}</td></tr>']
s += ['</tbody>']
s += ['</table>']
return '\n'.join(s)
@property
def has_incoherent(self):
""" Whether this sample contains the incoherent correlation functions or not. """
return False
@property
def has_currents(self):
""" Whether this sample contains the current correlation functions or not. """
return False
[docs]class StaticSample(Sample):
[docs] def to_dataframe(self):
""" Returns correlation functions as pandas dataframe """
df = pd.DataFrame()
for dim in self.dimensions:
df[dim] = self[dim].tolist() # to list to make q-points (N, 3) work in dataframe
for key in self.available_correlation_functions:
df[key] = self[key].reshape(-1, )
return df
[docs]class DynamicSample(Sample):
@property
def has_incoherent(self):
return 'Fqt_incoh' in self.available_correlation_functions
@property
def has_currents(self):
pair_string = '_'.join(self.pairs[0])
return f'Clqt_{pair_string}' in self.available_correlation_functions
[docs] def to_dataframe(self, q_index: int):
""" Returns correlation functions as pandas dataframe for the given q-index.
Parameters
----------
q_index
index of q-point to return
"""
df = pd.DataFrame()
for dim in self.dimensions:
if dim in ['q_points', 'q_norms']:
continue
df[dim] = self[dim]
for key in self.available_correlation_functions:
df[key] = self[key][q_index]
return df
[docs]def read_sample_from_npz(fname: str) -> Sample:
""" Read :class:`Sample <dynasor.sample.Sample>` from file.
Parameters
----------
fname
Path to the file (numpy npz format) from which to read
the :class:`Sample <dynasor.sample.Sample>` object.
"""
data_read = np.load(fname, allow_pickle=True)
data_dict = data_read['data_dict'].item()
meta_data = data_read['meta_data'].item()
if data_read['name'] == 'StaticSample':
return StaticSample(data_dict, **meta_data)
elif data_read['name'] == 'DynamicSample':
return DynamicSample(data_dict, **meta_data)
else:
return Sample(data_dict, **meta_data)