Source code for rtc_tools_meta.meta_mixin

import logging
from typing import Dict, List, Set, Tuple, Union

import casadi as ca
import numpy as np

from rtctools.optimization.optimization_problem import OptimizationProblem
from rtctools.optimization.timeseries import Timeseries

logger = logging.getLogger("rtctools")


class MetaMixin(OptimizationProblem):
    """Adds Compatibility with MetaOptimizationProblem models

    This mixin allows MetaOptimizationProblem models to assimilate your
    RTC-Tools model into its meta problem and impose constraints back
    onto your model. *Resistance is futile.*
    """

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # TODO: categorize fixed variables by run level?
        self.__fixed_interface_variables = {}
        self.__meta_solver_output = {}
        self.meta_level = None

    def _update_meta_level(self, meta_level: str):
        """Set meta run level.

        :param meta_level: String identifying meta run level.
        """
        logger.debug(
            f"MetaOptimizationProblem: {self.default_meta_level} activating meta level {meta_level}"
        )
        self.meta_level = meta_level

    @property
    def solver_output(self):
        if self.meta_level == self.default_meta_level:
            return super().solver_output
        else:
            try:
                return self.__meta_solver_output[self.meta_level]
            except KeyError:
                raise Exception(f"No results for meta level {self.meta_level}")

    def bounds(self):
        # Call parent class first for default values.
        bounds = super().bounds()
        for variable, timeseries in self.__fixed_interface_variables.items():
            bounds[variable] = (timeseries, timeseries)
        return bounds

    def _update_meta_solver_output(self, meta_solver_output: np.ndarray):
        """Record results from a meta optimization

        Takes the solver output from a meta optimization and redistributes it to
        the correct submodels.

        :param meta_solver_output: 1D numpy array of results calculated by
            ``self.optimize()``
        """
        self.__meta_solver_output[self.meta_level] = meta_solver_output

    def _fix_interface_variables(self, variables: Tuple[str]):
        """Fix interface variables of submodels

        Force subproblems to take the meta-optimal values of boundary conditions
        between submodels so that the submodels can be run independently

        :param variables: Tuple of variable names to be fixed
        """
        results = self.extract_results()
        for variable in variables:
            timeseries = Timeseries(self.times(variable), results[variable])
            self.__fixed_interface_variables[variable] = timeseries

    def optimize(self, *args, **kwargs):
        self._update_meta_level(self.default_meta_level)
        return super().optimize(*args, **kwargs)

    @property
    def default_meta_level(self):
        """Meta level corresponding to no meta modeling. Default is class name."""
        return self.__class__.__name__