Source code for pyretis.simulation.simulation

# -*- 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'])