Source code for pyretis.inout.formats.energy

# -*- coding: utf-8 -*-
# Copyright (c) 2023, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Module for formatting energy data from PyRETIS.

Important classes defined here
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

EnergyFormatter (:py:class:`.EnergyFormatter`)
    A class for formatting energy data from PyRETIS.

EnergyPathFormatter (:py:class:`.EnergyPathFormatter`)
    A class for formatting energy data for paths.

EnergyFile (:py:class:`.EnergyFile`)
    A class for handling PyRETIS energy files.

EnergyPathFile (:py:class:`.EnergyPathFile`)
    A class for handling PyRETIS energy path files.

"""
import logging
import numpy as np
from pyretis.inout.formats.formatter import OutputFormatter
from pyretis.inout.fileio import FileIO, read_some_lines
logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())

np.set_printoptions(legacy='1.25')

__all__ = [
    'EnergyFormatter',
    'EnergyPathFormatter',
    'EnergyFile',
    'EnergyPathFile',
]


[docs]class EnergyFormatter(OutputFormatter): """A class for formatting energy data from PyRETIS. This class handles formatting of energy data. The data is formatted in 5 columns: 1) Time, i.e. the step number. 2) Potential energy. 3) Kinetic energy. 4) Total energy, should equal the sum of the two previous columns. 5) Temperature. """ # Format for the energy files: ENERGY_FMT = ['{:>10d}'] + 5*['{:>14.6f}'] ENERGY_TERMS = ('vpot', 'ekin', 'etot', 'temp') HEADER = {'labels': ['Time', 'Potential', 'Kinetic', 'Total', 'Temperature'], 'width': [10, 14]}
[docs] def __init__(self, name='EnergyFormatter'): """Initialise the formatter for energy.""" super().__init__(name, header=self.HEADER)
[docs] def apply_format(self, step, energy): """Apply the energy format. Parameters ---------- step : int The current simulation step. energy : dict A dict with energy terms to format. Returns ------- out : string A string with the formatted energy data. """ towrite = [self.ENERGY_FMT[0].format(step)] for i, key in enumerate(self.ENERGY_TERMS): value = energy.get(key, None) if value is None: towrite.append(self.ENERGY_FMT[i + 1].format(float('nan'))) else: towrite.append(self.ENERGY_FMT[i + 1].format(float(value))) return ' '.join(towrite)
[docs] def format(self, step, data): """Yield formatted energy data. See :py:meth:.`apply_format`.""" yield self.apply_format(step, data)
[docs] def load(self, filename): """Load entire energy blocks into memory. Parameters ---------- filename : string The path/file name of the file we want to open. Yields ------ data_dict : dict This is the energy data read from the file, stored in a dict. This is for convenience so that each energy term can be accessed by `data_dict['data'][key]`. """ for blocks in read_some_lines(filename, line_parser=self.parse): data = np.array(blocks['data']) col = tuple(data.shape) col_max = min(col[1], len(self.ENERGY_TERMS) + 1) data_dict = {'comment': blocks['comment'], 'data': {'time': data[:, 0]}} for i in range(col_max-1): data_dict['data'][self.ENERGY_TERMS[i]] = data[:, i+1] yield data_dict
[docs]class EnergyPathFormatter(EnergyFormatter): """A class for formatting energy data for paths.""" ENERGY_TERMS = ('vpot', 'ekin') HEADER = {'labels': ['Time', 'Potential', 'Kinetic'], 'width': [10, 14]}
[docs] def __init__(self): """Initialise.""" super().__init__(name='EnergyPathFormatter') self.print_header = False
[docs] def format(self, step, data): """Format the order parameter data from a path. Parameters ---------- step : int The cycle number we are creating output for. data : tuple Here we assume that ``data[0]`` contains an object like :py:class:`.PathBase` and that ``data[1]`` is a string with the status for the path. Yields ------ out : string The strings to be written. """ path, status = data[0], data[1] if not path: # when nullmoves = False return move = path.generated yield f'# Cycle: {step}, status: {status}, move: {move}' yield self.header for i, phasepoint in enumerate(path.phasepoints): energy = {} for key in self.ENERGY_TERMS: energy[key] = getattr(phasepoint.particles, key, None) yield self.apply_format(i, energy)
[docs]class EnergyFile(FileIO): """A class for handling PyRETIS energy files."""
[docs] def __init__(self, filename, file_mode, backup=True): """Create the file object and attach the energy formatter.""" super().__init__(filename, file_mode, EnergyFormatter(), backup=backup)
[docs]class EnergyPathFile(FileIO): """A class for handling PyRETIS energy path files."""
[docs] def __init__(self, filename, file_mode, backup=True): """Create the file object and attach the energy formatter.""" super().__init__(filename, file_mode, EnergyPathFormatter(), backup=backup)