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.

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

    @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._data_keys:
            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)