Source code for pyretis.inout.formats.energy
# Copyright (c) 2026, 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.21')
__all__ = [
'EnergyFormatter',
'EnergyPathFormatter',
'EnergyFile',
'EnergyPathFile',
]
def _is_number(value):
"""Return True when ``value`` is a real (non-None, non-NaN) number."""
return value is not None and not (isinstance(value, float)
and np.isnan(value))
[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.
The path energy file carries the same four energy terms as the
per-step energy file -- potential, kinetic, total and temperature
(inherited :py:data:`EnergyFormatter.ENERGY_TERMS` /
:py:data:`EnergyFormatter.HEADER`). Engines that report only
potential and kinetic energies leave ``etot``/``temp`` unset, and
those columns are written as ``nan``. The reader
(:py:meth:`EnergyFormatter.load`) is column-count tolerant, so older
two-term path energy files still load.
"""
#: Cached ``dof*kB`` constant (``temp = 2*ekin/(dof*kB)``), shared
#: across formatter instances. Time reversal, swaps and restart reload
#: reuse frames without re-running the engine, and the derived
#: ``etot``/``temp`` are not serialised, so a reused frame can reach the
#: file with them unset. ``etot = vpot + ekin`` is reconstructed
#: directly; ``temp`` needs this per-system constant, which is recovered
#: exactly from any frame that carries both ``ekin`` and ``temp`` and is
#: a true simulation-global, so caching it lets even a fully-reused path
#: (no temperature-bearing frame of its own) reconstruct temperature.
_dof_kb = None
[docs]
def __init__(self):
"""Initialise."""
super().__init__(name='EnergyPathFormatter')
self.print_header = False
[docs]
def _reconstruct(self, energies):
"""Fill missing total energy / temperature from the stored vpot/ekin.
Reconstructs the derived terms so the path energy file stays
reproducible across frame reuse and restart (see :attr:`_dof_kb`).
Only genuinely-missing terms are filled; engine-reported values are
left untouched.
"""
for energy in energies:
ekin, temp = energy.get('ekin'), energy.get('temp')
if _is_number(ekin) and _is_number(temp) and temp != 0.0:
EnergyPathFormatter._dof_kb = 2.0 * ekin / temp
break
dof_kb = EnergyPathFormatter._dof_kb
for energy in energies:
vpot, ekin = energy.get('vpot'), energy.get('ekin')
if (not _is_number(energy.get('etot'))
and _is_number(vpot) and _is_number(ekin)):
energy['etot'] = vpot + ekin
if (not _is_number(energy.get('temp'))
and _is_number(ekin) and dof_kb is not None):
energy['temp'] = 2.0 * ekin / dof_kb
[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
energies = []
for phasepoint in path.phasepoints:
energies.append({key: getattr(phasepoint.particles, key, None)
for key in self.ENERGY_TERMS})
self._reconstruct(energies)
for i, energy in enumerate(energies):
yield self.apply_format(i, energy)