Source code for dynasor.sample

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.

        Dictionary with correlation functions.
        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'] """
            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.

            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)

    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))

    def dimensions(self):
        r"""The dimensions for the samples, e.g., for
        :math:`S(q, \omega)` the dimensions of 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))

    def atom_types(self):
        if self._atom_types is None:
            return None
        return self._atom_types.copy()

    def particle_counts(self):
        if self._particle_counts is None:
            return None
        return self._particle_counts.copy()

    def pairs(self):
        if self._pairs is None:
            return None
        return self._pairs.copy()

    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:')
        for key in self._data_keys:
            s_i = f'{key:15} with shape: {np.shape(getattr(self,key))}'
        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>'
        s += ['<tbody>']
        s += ['<tr><td style="text-align: left">Atom types</td>'
        s += ['<tr><td style="text-align: left">Pairs</td>'
        s += ['<tr><td style="text-align: left">Particle counts</td>'
        s += ['<tr><td style="text-align: left">Simulations cell</td>'
        for key in self._data_keys:
            s += [f'<tr><td style="text-align: left">{key}</td>'
        s += ['</tbody>']
        s += ['</table>']
        return '\n'.join(s)

    def has_incoherent(self):
        """ Whether this sample contains the incoherent correlation functions or not. """
        return False

    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)