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)