# -*- coding: utf-8 -*-
# Copyright (c) 2023, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Definitions of generic simulation objects.
This module defines the generic simulation object. This is the base
class for all other simulations.
Important classes defined here
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Simulation (:py:class:`.Simulation`)
A class defining a generic simulation.
"""
import logging
import copy
import os
from pyretis.simulation.simulation_task import SimulationTask
from pyretis.inout.screen import print_to_screen
from pyretis.inout.simulationio import task_from_settings
from pyretis.inout.restart import write_restart_file
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())
__all__ = ['Simulation']
[docs]class Simulation:
"""This class defines a generic simulation.
Attributes
----------
cycle : dict of integers
This dictionary stores information about the number of cycles.
The keywords are:
* `step`: The current cycle number.
* `startcycle`: The cycle number we started at.
* `endcycle`: Represents the cycle number where the simulation
should end.
* `stepno`: The number of cycles we have performed to arrive at
cycle number given by `cycle['step']`. Note that `cycle['stepno']`
might be different from `cycle['step']` since `cycle['startcycle']`
might be != 0.
exe_dir : string
The path we are running the simulation from.
restart_freq : integer
The frequency for creating restart files.
first_step : boolean
True if the first step has not been executed yet.
system : object like :py:class:`.System`
This is the system the simulation will act on.
simulation_output : list of dicts
This list defines the output tasks associated
with the simulation.
simulation_type : string
An identifier for the simulation.
tasks : list of objects like :py:class:`.SimulationTask`
This is the list of simulation tasks to execute.
"""
simulation_type = 'generic'
simulation_output = []
[docs] def __init__(self, settings, controls):
"""Initialise the simulation object.
Parameters
----------
controls: dict of parameters to set up and control the simulations.
It contains:
* `steps`: int, optional
The number of simulation steps to perform.
* `startcycle`: int, optional
The cycle we start the simulation on, useful for restarts.
* `endcycle`: int, optional
The cycle we end the simulation to, useful in restarts.
* `rgen`: object like :py:class:`.RandomGenerator`
The random generator that will be used for the
paths that required random numbers.
settings : dict
Contains all the simulation settings.
"""
steps = controls.get('steps', 0)
startcycle = controls.get('startcycle', 0)
end = controls.get('endcycle', steps)
self.cycle = {'step': startcycle, 'endcycle': end,
'startcycle': startcycle, 'stepno': 0, 'steps': steps}
self.tasks = []
self.output_tasks = []
self.first_step = True
self.system = None
self.restart_freq = None
self.exe_dir = None
self.settings = settings
[docs] def extend_cycles(self, steps):
"""Extend a simulation with the given number of steps.
Parameters
----------
steps : int
The number of steps to extend the simulation with.
Returns
-------
out : None
Returns `None` but modifies `self.cycle`.
"""
self.cycle['startcycle'] = self.cycle['stepno']
self.cycle['endcycle'] = self.cycle['startcycle'] + steps
[docs] def is_finished(self):
"""Determine if the simulation is finished.
In this object, the simulation is done if the current step
number is larger than the end cycle. Note that the number of
steps performed is dependent on the value of
`self.cycle['startcycle']`.
Returns
-------
out : boolean
True if the simulation is finished, False otherwise.
"""
return self.cycle['step'] >= self.cycle['endcycle']
[docs] def step(self):
"""Execute a simulation step.
Here, the tasks in :py:attr:`.tasks` will be executed
sequentially.
Returns
-------
out : dict
This dictionary contains the results of the defined tasks.
Note
----
This function will have 'side effects' and update/change
the state of other attached variables such as the system or
other variables that are not explicitly shown. This is intended
and the behavior is defined by the tasks in
:py:attr:`.tasks`.
"""
if not self.first_step:
self.cycle['step'] += 1
self.cycle['stepno'] += 1
results = self.execute_tasks()
if self.first_step:
self.first_step = False
return results
[docs] def execute_tasks(self):
"""Execute all the tasks in sequential order.
Returns
-------
results : dict
The results from the different tasks (if any).
"""
results = {'cycle': self.cycle.copy()}
for task in self.tasks:
if not self.first_step or task.run_first():
resi = task.execute(self.cycle)
if task.result is not None:
results[task.result] = resi
results['system'] = self.system
return results
[docs] def add_task(self, task, position=None):
"""Add a new simulation task.
A task can still be added manually by simply appending to
py:attr:`.tasks`. This function will, however, do some
checks so that the task added can be executed.
Parameters
----------
task : dict
A dict defining the task. A task is represented by an
object of type :py:class:`.SimulationTask` with some
additional settings on how to store the output
and when to execute the task.
The keywords in the dict defining the task are:
* `func`: Callable, this is a function to execute in the
task.
* `args`: List, with arguments for the function.
* `kwargs`: Dict, with the keyword arguments for the
function.
* `when`: Dict, which defines when the task should be
executed.
* `first`: Boolean, determines if the task should be
executed on the initial step, i.e. before the
full simulation starts.
* `result`: String, used to label the result.
position : int, optional
Can be used to give the tasks a specific position in the
task list.
"""
try:
new_task = SimulationTask(task['func'],
args=task.get('args', None),
kwargs=task.get('kwargs', None),
when=task.get('when', None),
result=task.get('result', None),
first=task.get('first', False))
if position is None:
self.tasks.append(new_task)
else:
self.tasks.insert(position, new_task)
return True
except AssertionError:
logger.warning('Could not add task: %s', task)
return False
[docs] def run(self):
"""Run a simulation.
The intended usage is for simulations where all tasks have
been defined in :py:attr:`self.tasks`.
Note
----
This function will just run the tasks via executing
:py:meth:`.step` In general, this is probably too generic for
the simulation you want, if you are creating a custom simulation.
Please consider customizing the :py:meth:`.run` (or the
:py:meth:`.step`) method of your simulation class.
Yields
------
out : dict
This dictionary contains the results from the simulation.
"""
while not self.is_finished():
result = self.step()
for task in self.output_tasks:
task.output(result)
self.write_restart()
if self.soft_exit():
yield result
break
yield result
[docs] def __str__(self):
"""Just a small function to return some info about the simulation."""
ntask = len(self.tasks)
mtask = 'task' if ntask == 1 else 'tasks'
msg = [f'Generic simulation with {ntask} {mtask}.']
for i, task in enumerate(self.tasks):
msg += [f'* Task no. {i}']
for j, line in enumerate(str(task).split('\n')):
if j > 0:
msg += [line]
return '\n'.join(msg)
[docs] def set_up_output(self, settings, progress=False):
"""Set up output from the simulation.
This includes the predefined output tasks, but also output
related to the restart file(s).
Parameters
----------
settings : dict
These are the simulation settings.
progress : boolean
For some simulations, the user may select to display a
progress bar, we then need to disable the screen output.
"""
logging.debug('Setting up output for simulation %s',
self.__class__.__name__)
# Create the output tasks:
self.create_output_tasks(settings, progress=progress)
# Do set-up for restart output:
self.restart_freq = settings['output'].get('restart-file', -1)
if self.restart_freq < 1:
self.restart_freq = None
logger.warning('Writing of restart file(s) disabled!')
logger.debug('Setting restart frequency for simulation %s',
self.restart_freq)
self.exe_dir = settings['simulation'].get('exe_path', os.getcwd())
[docs] def create_output_tasks(self, settings, progress=False):
"""Create output tasks for the simulation.
This method will generate output tasks based on the tasks
listed in :py:attr:`.simulation_output`.
Parameters
----------
settings : dict
These are the simulation settings.
progress : boolean
For some simulations, the user may select to display a
progress bar, we then need to disable the screen output.
"""
logging.debug('Clearing output tasks & adding pre-defined ones')
engine = getattr(self, 'engine', None)
order_function = getattr(self, 'order_function', None)
self.output_tasks = []
directory = settings['simulation'].get('exe_path', None)
for task_dict in self.simulation_output:
if 'order' in task_dict['type'] and order_function is None:
continue
task = task_from_settings(task_dict, settings, directory, engine,
progress)
if task is not None:
logger.debug('Created output task:\n%s', task)
self.output_tasks.append(task)
[docs] def soft_exit(self):
"""Force simulation to stop at the current step."""
exit_file = 'EXIT'
if self.exe_dir:
exit_file = os.path.join(self.exe_dir, exit_file)
if os.path.isfile(exit_file):
logger.info('Exit file found - will do a soft exit.')
print_to_screen('Exit file found - will do a soft exit.',
level='warning')
# Write restart file...
self.write_restart(now=True)
# Close output files...
for task in self.output_tasks:
if task.target == 'file':
task.writer.close()
return True
return False
[docs] def write_restart(self, now=False):
"""Create a restart file.
Parameters
----------
now : boolean, optional
If True, the output file will be written irrespective of the
step number.
"""
if now or (self.restart_freq is not None and
self.cycle['step'] % self.restart_freq == 0):
out = 'pyretis.restart'
if self.exe_dir:
out = os.path.join(self.exe_dir, out)
write_restart_file(out, self)
[docs] def restart_info(self):
"""Return information which can be used to restart the simulation.
Returns
-------
info : dict,
Contains all the updated simulation settings and counters.
"""
info = {}
info['settings'] = self.settings
for key in ('simulation', 'system', 'particles'):
if key not in info:
info[key] = {}
info['simulation']['restart'] = 'pyretis.restart'
info['simulation']['cycle'] = copy.deepcopy(self.cycle)
info['simulation']['type'] = self.simulation_type
if hasattr(self, 'system') and self.system is not None:
info['system'].update(self.system.restart_info())
if hasattr(self.system, 'particles') and \
self.system.particles is not None:
info['particles'].update(self.system.particles.restart_info())
return info
[docs] def load_restart_info(self, info):
"""Load restart information.
Note, we do not change the ``end`` property here as we probably
are extending a simulation.
Parameters
----------
info : dict
The dictionary with the restart information.
"""
for key, val in info['simulation']['cycle'].items():
if key in {'steps', 'endcycle'}:
self.cycle['startcycle'] = copy.deepcopy(val)
else:
self.cycle[key] = copy.deepcopy(val)
if info.get('system') is not None and\
hasattr(self, 'system') and self.system is not None:
self.system.load_restart_info(info['system'])