Source code for aiida_optimize.engines.base

# -*- coding: utf-8 -*-

# © 2017-2019, ETH Zurich, Institut für Theoretische Physik
# Author: Dominik Gresch <greschd@gmx.ch>
"""
Defines the base classes needed to implement an optimization engine.
"""

from __future__ import annotations

from abc import ABCMeta, abstractmethod
import typing as ty

import yaml

from ._result_mapping import Result, ResultMapping

yaml.representer.Representer.add_representer(ABCMeta, yaml.representer.Representer.represent_name)  # type: ignore

__all__ = (
    "OptimizationEngineImpl",
    "OptimizationEngineImplWithOutputs",
    "OptimizationEngineWrapper",
)


[docs]class OptimizationEngineImpl: """ Base class for the stateful optimization engine implementation. """ __metaclass__ = ABCMeta def __init__( self, logger: ty.Any, result_state: ty.Optional[ty.Dict[int, Result]] = None ) -> None: self._logger = logger self._result_mapping = ResultMapping.from_state(result_state)
[docs] @classmethod def from_state(cls, state: ty.Dict[str, ty.Any]) -> OptimizationEngineImpl: """ Create an instance of the class from the serialized state. """ return cls(**state)
@property def state(self) -> ty.Dict[str, ty.Any]: """ The serialized state of the instance, including the result mapping. """ return dict(result_state=self._result_mapping.state, **self._state) @property @abstractmethod def _state(self) -> ty.Dict[str, ty.Any]: """ The serialized state of the instance, without the result mapping. This function needs to be implemented by child classes. """ @property @abstractmethod def is_finished(self) -> bool: """ Returns true if the optimization is finished. """ @property def is_finished_ok(self) -> bool: """ Returns true if the optimization is finished without error. """ return self.is_finished
[docs] def create_inputs(self): """ Creates the inputs and adds them to the result mapping. """ return self._result_mapping.add_inputs(self._create_inputs())
[docs] @abstractmethod def _create_inputs(self): """ Creates the inputs for evaluations that need to be launched. This function needs to be implemented by child classes. """
[docs] def update(self, outputs) -> None: """ Updates the result mapping and engine instance with the evaluation outputs. """ self._result_mapping.add_outputs(outputs) self._update(outputs)
[docs] @abstractmethod def _update(self, outputs: ty.Dict[int, ty.Any]) -> None: """ Updates the engine instance with the evaluation outputs. This method needs to be implemented by child classes. """
@property def result_index(self) -> int: """ Returns the index (in the result mapping) of the optimal evaluation. """ index, _, _ = self._get_optimal_result() return index @property def result_input_value(self) -> ty.Any: """ Return the input value of the optimal evaluation. """ _, value, _ = self._get_optimal_result() return value @property def result_output_value(self) -> ty.Any: """ Return the output value of the optimal evaluation. """ _, _, value = self._get_optimal_result() return value
[docs] @abstractmethod def _get_optimal_result(self) -> ty.Tuple[int, ty.Any, ty.Any]: """ Return the index, input value, and output value of the best evaluation process. """
[docs]class OptimizationEngineImplWithOutputs: """ Base class for the optimization engine implementation emitting custom outputs. """
[docs] @abstractmethod def get_engine_outputs(self) -> ty.Dict[str, ty.Any]: """ Return the custom outputs created by the engine, at the end of the run. The result must be compatible with `WorkChain.out`, i.e. its keys are labels, and the values are either AiiDA nodes, or nested dictionaries. All AiiDA nodes returned *must* already be stored. """
[docs]class OptimizationEngineWrapper: """ Base class for wrappers that supply the public interface for optimization engines. """ __metaclass__ = ABCMeta _IMPL_CLASS: ty.Type[OptimizationEngineImpl] def __new__(cls, *args, **kwargs): return cls._IMPL_CLASS(*args, **kwargs) @classmethod def from_state(cls, state, logger): return cls._IMPL_CLASS(logger=logger, **state)