gemseo-multi-fidelity 0.0.1__py3-none-any.whl

gemseo_multi_fidelity/core/scenario_adapter.py

@@ -0,0 +1,568 @@
+# Copyright 2021 IRT Saint Exupéry, https://www.irt-saintexupery.com
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License version 3 as published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this program; if not, write to the Free Software Foundation,
+# Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
+
+# Copyright (c) 2019 AIRBUS OPERATIONS
+
+#
+# Contributors:
+#    INITIAL AUTHORS - API and implementation and/or documentation
+#      :author: Romain Olivanti
+#    OTHER AUTHORS   - MACROSCOPIC CHANGES
+#        Francois Gallard: portage to GEMSEO 6
+"""Multi-fidelity scenario adapter."""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+from typing import Any
+
+from gemseo.algos.database import Database
+from gemseo.core.discipline.discipline import Discipline
+from gemseo.scenarios.mdo_scenario import MDOScenario
+
+from gemseo_multi_fidelity.core.corr_function import MDOCorrectedFunction
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from collections.abc import Mapping
+
+    from gemseo.core._process_flow.base_process_flow import BaseProcessFlow
+    from gemseo.core.mdo_functions.mdo_function import MDOFunction
+    from gemseo.typing import StrKeyMapping
+    from numpy.typing import NDArray
+
+LOGGER = logging.getLogger(__name__)
+
+
+class CallbackReturnError(Exception):
+    """Error to raise if a return is requested at runtime."""
+
+    def __init__(self, msg: str | None = None) -> None:
+        """Constructor.
+
+        Args:
+            msg: explanation for the error.
+        """
+        super().__init__()
+        self.msg = msg
+
+    def __str__(self) -> str:
+        """Custom string representation."""
+        return repr(self.msg)
+
+
+class MFScenarioAdapter(Discipline):
+    """An adapter class to run a scenario in a multi-fidelity context."""
+
+    X_START = "x_start"
+    X_BEST = "x_best"
+
+    INPUT_OPTIONS = "input_options"
+    OUTPUT_OPTIONS = "output_options"
+
+    X_REF = "x_reference"
+    OUTPUT_REF = "output_reference"
+
+    ON_ENTRY_CALLBACKS = "on_entry_callbacks"
+    RETURN_CALLBACK = "return_callback"
+    ON_EXIT_CALLBACKS = "on_exit_callbacks"
+
+    OUTPUT_STATUS = "output_status"
+    NOT_STARTED = "not started"
+    EARLY_RETURN = "asked to return"
+    RUN_COMPLETED = "completed"
+    N_ITER_RUN = "n_iters"
+
+    MAX_CORR_RAD = "max_corr_rad"
+    HOTSTART_OPTIONS = "hotstart_options"
+    HESS_PAIRS = "input_hess_pairs"
+    PRECOND = "input_precond"
+
+    DEF_HOTSTART_OPTIONS_GETTER = "get_hotstart_options"
+
+    scenario: MDOScenario
+    """The wrapped scenario"""
+
+    auto_detect_grammar_files = True
+
+    def __init__(self, scenario: MDOScenario) -> None:
+        """Constructor.
+
+        Args:
+            scenario: The scenario to run.
+        """
+        if not isinstance(scenario, MDOScenario):
+            msg = "input scenario must be an MDOScenario"
+            raise TypeError(msg)
+
+        super().__init__(name=scenario.name)
+
+        self.scenario = scenario
+
+        self._opt_status = self.NOT_STARTED
+        self._return_callback: callable | None = None
+        self._hotstart_options_set: bool = False
+        self._init_return_callback()
+
+    def _init_return_callback(self):
+        """Initialize return callback.
+
+        Initializes a return callback linked to the database of the optimization
+        problem. By default, the callback has no effect.
+        """
+
+        def must_return(x_vect=None):
+            return False
+
+        self._return_callback = must_return
+        self.scenario.formulation.optimization_problem.database.add_new_iter_listener(
+            self._return_callback
+        )
+
+    @staticmethod
+    def check_callback(callback: Callable) -> None:
+        """Check if the callback provided is valid.
+
+        Args:
+            callback: The callback to check.
+        """
+        if not callable(callback):
+            msg = f"{callback} must be callable"
+            raise TypeError(msg)
+
+    @staticmethod
+    def _check_return_callback(callback: Callable) -> None:
+        """Check return callback.
+
+        Checks that the callback provided is a valid return callback, i.e. that it only
+        returns a bool as output and that it should not already order to return at check
+        time.
+
+        Args:
+            callback: The callback to check.
+        """
+        MFScenarioAdapter.check_callback(callback)
+
+        # Specific checks
+        return_flag = callback()
+        if not isinstance(return_flag, bool):
+            msg = "return_callback must return a bool"
+            raise TypeError(msg)
+
+        if return_flag:
+            LOGGER.warning(
+                "return_callback is asking to return before the "
+                "beginning of the optimization"
+            )
+
+    def _link_return_callback(self, return_callback: Callable) -> None:
+        """Link return callback.
+
+        Implements the return callback using a listener linked to the optimization
+        database of the scenario. ``return_callback`` must return a bool and take the
+        design vector as argument.
+        At least two stores are allowed to be performed irrespectively to the return
+        callback in order to avoid deadlocks during optimization (overrides the return
+        flag).
+
+        Args:
+            return_callback: The callback allowing an early return.
+        """
+        self._check_return_callback(return_callback)
+
+        # Delete the current listener
+        opt_prob = self.scenario.formulation.optimization_problem
+        # TODO access should be cleaner than this, add a method in GEMSEO
+        new_it_listeners = opt_prob.database._Database__new_iter_listeners
+        with contextlib.suppress(ValueError):
+            new_it_listeners.remove(self._return_callback)
+        # Make sure at least one real optimization iteration can be performed by the
+        # MDOScenario, build a specific callback.
+        it_init = len(opt_prob.database)
+        it_return_allowed = it_init + 1
+
+        def must_return(x_vect=None) -> None:
+            """Callback making sure that at least on real optimization is performed."""
+            asked_to_return = return_callback()
+            curr_it = len(opt_prob.database)
+            if asked_to_return and (curr_it > it_return_allowed):
+                msg = "Ordered to return by the callback"
+                raise CallbackReturnError(msg)
+
+        # Add the new callback
+        self._return_callback = must_return
+        opt_prob.database.add_new_iter_listener(self._return_callback)
+
+    def run_callbacks_from_local_data(self, callback_type: str) -> list:
+        """Run the specified callback type from the local data.
+
+        Also checks that the callbacks defined in the default_inputs are also called if
+        different from the ones defined in the local_data.
+        This is explained by the fact that the behavior expected here is not to override
+        but to append.
+
+        Args:
+            callback_type: The key in the local data.
+
+        Returns:
+            The list ??? # TODO
+        """
+        callbacks = self.io.data.get(callback_type)
+        out = []
+
+        if callbacks is not None:
+            # Make sure that the callbacks defined in the default_inputs are
+            # not simply overridden
+            if callable(callbacks):
+                # Make sure it is a list to ease the check
+                callbacks = [callbacks]
+            # Add them if there are not there yet
+            callbacks_in_default = self.default_input_data.get(callback_type, [])
+            if callable(callbacks_in_default):
+                callbacks_in_default = [callbacks_in_default]
+            for callback in callbacks_in_default:
+                if callback not in callbacks:
+                    callbacks.append(callback)
+            out = self.run_callbacks(callbacks, self.io.data)
+        return out
+
+    @staticmethod
+    def run_callbacks(
+        callbacks: Iterable[Callable] | Callable,
+        *args: list[Any],
+        **kwargs: Mapping[str, Any],
+    ) -> list[Callable]:
+        """Run the specified callbacks with the same args sequentially.
+
+        Args:
+            callbacks: The callbacks to run.
+            args: The args.
+            kwargs: The keyword arguments.
+
+        Returns:
+            The list of outputs.
+        """
+        if isinstance(callbacks, Iterable):
+            return [callback(*args, **kwargs) for callback in callbacks]
+
+        return [callbacks(*args, **kwargs)]
+
+    def get_all_nonproc_funcs(self) -> list[MDOFunction]:
+        """Get all non-processed functions.
+
+        Convenience method to get all the non-processed functions even when the
+        functions have not been preprocessed yet.
+
+        Returns:
+            The list of the non-processed functions.
+        """
+        opt_prob = self.scenario.formulation.optimization_problem
+
+        if opt_prob._functions_are_preprocessed:
+            return opt_prob.functions
+        return opt_prob.original_functions
+
+    def get_algo_lib(self) -> str:
+        """Convenience method to get the algo lib specified to run the scenario.
+
+        Returns:
+            The lib.
+        """
+        # The algo name is either set by the input data or already defined in the
+        # default inputs of the scenario
+        tag_algo = "algo_name"
+        default_algo = None
+        if self.scenario._settings is not None:
+            default_algo = self.scenario._settings.algo_name
+        algo_name = self.io.data.get(tag_algo, default_algo)
+        if algo_name is None:
+            msg = "No algorithm specified to execute the scenario"
+            raise RuntimeError(msg)
+
+        return self.scenario._algo_factory.create(algo_name)
+
+    def _set_func_reference(
+        self, func: MDOCorrectedFunction, x_ref: NDArray, output_ref: dict
+    ) -> None:
+        """Set function reference.
+
+        Sets the reference data for the specified function.
+        Nothing is done if no data are provided for the specified function.
+
+        Args:
+            func: The MDOCorrectedFunction to update.
+            x_ref: The reference point.
+            output_ref: The reference data to set.
+        """
+        max_norm_corr = self.io.data.get(self.MAX_CORR_RAD)
+        func_name = func.name
+
+        val_ref = output_ref.get(func_name)
+        if val_ref is not None:
+            grad_ref = output_ref.get(Database.GRAD_TAG + func_name)
+            if func.original is not None and isinstance(
+                func.original, MDOCorrectedFunction
+            ):
+                func = func.original
+            func.set_reference(
+                x_ref, val_ref, grad_ref=grad_ref, max_norm_thr=max_norm_corr
+            )
+
+    def _set_reference(self) -> None:
+        """Set the reference data provided."""
+        x_ref = self.io.data.get(self.X_REF)
+        output_ref = self.io.data.get(self.OUTPUT_REF)
+
+        if x_ref is None or output_ref is None:
+            # Nothing to do
+            return
+
+        opt_prob = self.scenario.formulation.optimization_problem
+
+        # Correction will be set, therefore, clear the current optimization data as it
+        # may no longer be consistent
+        opt_prob.database.clear()
+
+        to_array = self.scenario.design_space.convert_dict_to_array
+        # Convert to arrays
+        if isinstance(x_ref, dict):
+            x_ref = to_array(x_ref)
+        for func, jac in output_ref.items():
+            if isinstance(jac, dict):
+                output_ref[func] = to_array(x_ref)
+
+        # Apply the reference data
+        for func in self.get_all_nonproc_funcs():
+            self._set_func_reference(func, x_ref, output_ref)
+
+    def _filter_subprocess_inputs(self) -> dict:
+        """Filter subprocess inputs.
+
+        Filters self.local_data to keep only the data required to run the MDOScenario.
+
+        Returns:
+            The filtered inputs.
+        """
+        LOGGER.warning(
+            "ScenarioAdapter._filter_subprocess_inputs: Deactivated subprocess inputs!"
+        )
+        return {}
+        input_data = {}
+
+        key_algo = "algo_name"
+        key_algo_opt = "algo_options"
+        key_max_iter = "max_iter"
+
+        # Check if the algo specified in self.local_data is the same as the one defined
+        # in the default inputs
+        def_algo = (
+            self.scenario._settings.algo_name
+            if self.scenario._settings is not None
+            else None
+        )
+        algo = self.io.data.get(key_algo)
+        #
+        new_algo_options = self.io.data.get(key_algo_opt, {})
+        #
+        if def_algo is not None and (algo == def_algo or algo is None):
+            # Same algo Try to load the default inputs in the input data as only a
+            # limited number of options may have been updated
+            algo_options = self.scenario._settings.model_dump()
+            algo_options.update(new_algo_options)
+        else:
+            # The new algo differs from the default one
+            # Use only new options
+            algo_options = new_algo_options
+
+        max_iter = self.io.data.get(key_max_iter)
+
+        # Handle preconditioner and hess approx pairs inputs
+        other_algo_options = self._proccess_lib_options()
+        algo_options.update(other_algo_options)
+
+        input_data = {
+            key: data
+            for key, data in zip(
+                [key_algo, key_max_iter], [algo, max_iter], strict=False
+            )
+            if data is not None
+        }
+
+        input_data.update(algo_options)
+
+        return input_data
+
+    def _proccess_lib_options(self) -> dict:
+        """Process lib-specific options.
+
+        Returns:
+            The processed options.
+        """
+        lib_options = {}
+        # Preconditioner
+        precond = self.io.data.get(self.PRECOND)
+        algo_lib = self.get_algo_lib()
+        if precond is not None:
+            try:
+                precond_lib_name = algo_lib.PRECOND
+                lib_options[precond_lib_name] = precond
+            except AttributeError:
+                pass
+        # Hessian approx pairs
+        pairs = self.io.data.get(self.HESS_PAIRS)
+        if pairs is not None:
+            try:
+                pairs_lib_name = algo_lib.HESS_PAIRS
+                lib_options[pairs_lib_name] = pairs
+            except AttributeError:
+                pass
+        return lib_options
+
+    def _set_inputs(self) -> dict:
+        """Tasks to perform before executing self.scenario.
+
+        Returns:
+            The input data to run the MDOScenario.
+        """
+        self._opt_status = self.NOT_STARTED
+        x_start = self.io.data[self.X_START]
+        self.scenario.design_space.set_current_value(x_start)
+
+        # Unpack forwarded options
+        input_options = self.io.data.get(self.INPUT_OPTIONS)
+        if input_options is not None:
+            for option, value in input_options.items():
+                self.io.data[option] = value
+
+        # Set the reference data
+        self._set_reference()
+
+        # Link the update callback
+        return_callback = self.io.data.get(self.RETURN_CALLBACK)
+        if return_callback is not None:
+            self._link_return_callback(return_callback)
+
+        # Filter scenario inputs
+        return self._filter_subprocess_inputs()
+
+    def _set_outputs(self) -> None:
+        """Store the outputs from the optimization."""
+        # Output the best data expressed in the upper design space
+        opt_prob = self.scenario.formulation.optimization_problem
+
+        # Get the best point
+        x_best = opt_prob.optimum[1]
+        x_best = opt_prob.design_space.convert_array_to_dict(x_best)
+
+        # Update the output data
+        self.io.data[self.X_BEST] = x_best
+
+        # Update the output status
+        self.io.data[self.OUTPUT_STATUS] = self._opt_status
+
+        # Set the potential hotstart options in the event of a restart
+        # At this level
+        self._set_hotstart_options()
+
+    def _set_hotstart_options(self) -> None:
+        """Set hotstart options.
+
+        Either tries to recover hotstart options from the lib itself, using a specific
+        getter if a string is provided in self.local_data[self.HOTSTART_OPTIONS] or uses
+        the options provided if a dict is provided.
+        If nothing is provided, a default getter is called.
+        If the latter is not recognized by the lib, nothing is done.
+        """
+        hotstart_options = self.io.data.get(
+            self.HOTSTART_OPTIONS, self.DEF_HOTSTART_OPTIONS_GETTER
+        )
+        options = None
+        self._hotstart_options_set = False
+        if isinstance(hotstart_options, dict):
+            # Simply use the options provided
+            options = hotstart_options
+        elif isinstance(hotstart_options, str):
+            # Try to recover the options from the specified getter
+            algo_lib = self.get_algo_lib()
+            try:
+                hotstart_options_getter = getattr(algo_lib, hotstart_options)
+                opt_prob = self.scenario.formulation.optimization_problem
+                options = hotstart_options_getter(
+                    opt_prob.database, opt_prob.objective.name
+                )
+            except:  # noqa: E722
+                # Nothing can be done
+                pass
+
+        if options is not None:
+            # Set the options
+            try:
+                algo_options = self.io.data["algo_options"]
+                algo_options.update(options)
+            except KeyError:
+                # No algo options have defined yet
+                self.io.data["algo_options"] = options
+            self._hotstart_options_set = True
+
+    def _run_subprocess(self, input_data: StrKeyMapping) -> None:
+        """Run the subprocess, here the optimization of the scenario.
+
+        Args:
+            input_data: The input data.
+        """
+        self.scenario.execute(**input_data)
+
+    def _run(self, input_data: StrKeyMapping) -> None:
+        """Run.
+
+        Processes the input data potentially:
+         - setting the reference data
+         - setting the return callback
+        then runs:
+         - the on entry callbacks
+         - the scenario
+         - the on exit callbacks.
+
+        Args:
+             input_data: The input data.
+        """
+        # Run the on entry callback(s)
+        self.run_callbacks_from_local_data(self.ON_ENTRY_CALLBACKS)
+
+        inputs_subprocess = self._set_inputs()
+        # Run the subprocess
+        try:
+            self._run_subprocess(inputs_subprocess)
+            self._opt_status = self.RUN_COMPLETED
+        except CallbackReturnError:
+            # Asked to return by the callback
+            self._opt_status = self.EARLY_RETURN
+
+        # Set the output data
+        self._set_outputs()
+
+        # Run the on exit callback(s)
+        self.run_callbacks_from_local_data(self.ON_EXIT_CALLBACKS)
+
+    def get_process_flow(self) -> BaseProcessFlow:
+        """Get the process flow.
+
+        Returns:
+            The process flow.
+        """
+        return self.scenario.get_process_flow()