A Sellar problem with custom data converters

from __future__ import annotations

import operator
from math import exp
from typing import TYPE_CHECKING

from numpy import array
from numpy import ones

from gemseo import create_scenario
from gemseo import set_data_converters
from gemseo.algos.design_space import DesignSpace
from gemseo.core.discipline import Discipline

if TYPE_CHECKING:
    from gemseo.typing import StrKeyMapping

As compared to the original Sellar example, here the y_1 variable is a dictionary. and the y_2 variable is a 2D array The grammars are changed to SimpleGrammar for simplicity, otherwise, with the default JSONGrammar, more complex definitions would be required.

DEFAULT_INPUT_DATA = {
    "x": ones(1),
    "z": array([1.0, 0.0]),
    "y_1": {"dummy": ones(1)},
    "y_2": ones((1, 1)),
}


class SellarSystem(Discipline):
    default_grammar_type = Discipline.default_grammar_type.SIMPLE

    def __init__(self) -> None:
        super().__init__()
        default_input_data = DEFAULT_INPUT_DATA.copy()
        self.input_grammar.update_from_data(default_input_data)
        self.output_grammar.update_from_names(["obj", "c_1", "c_2"])
        self.default_input_data = default_input_data

    def _run(self, input_data: StrKeyMapping) -> StrKeyMapping | None:
        x = input_data["x"]
        z = input_data["z"]
        y_1 = input_data["y_1"]["dummy"]
        y_2 = input_data["y_2"][0]
        return {
            "obj": array([x[0] ** 2 + z[1] + y_1[0] ** 2 + exp(-y_2[0])]),
            "c_1": array([3.16 - y_1[0] ** 2]),
            "c_2": array([y_2[0] - 24.0]),
        }


class Sellar1(Discipline):
    default_grammar_type = Discipline.default_grammar_type.SIMPLE

    def __init__(self) -> None:
        super().__init__()
        default_input_data = DEFAULT_INPUT_DATA.copy()
        # Keep the y_1 data for easily defining the output grammar.
        output_grammar_data = {"y_1": default_input_data.pop("y_1")}
        self.input_grammar.update_from_data(default_input_data)
        self.output_grammar.update_from_data(output_grammar_data)
        self.default_input_data = default_input_data

    def _run(self, input_data: StrKeyMapping) -> StrKeyMapping | None:
        x = input_data["x"]
        z = input_data["z"]
        y_2 = input_data["y_2"][0]
        return {
            "y_1": {"dummy": array([(z[0] ** 2 + z[1] + x[0] - 0.2 * y_2[0]) ** 0.5])}
        }


class Sellar2(Discipline):
    default_grammar_type = Discipline.default_grammar_type.SIMPLE

    def __init__(self) -> None:
        super().__init__()
        default_input_data = DEFAULT_INPUT_DATA.copy()
        del default_input_data["y_2"]
        del default_input_data["x"]
        self.input_grammar.update_from_data(default_input_data)
        self.output_grammar.update_from_names(["y_2"])
        self.default_input_data = default_input_data

    def _run(self, input_data: StrKeyMapping) -> StrKeyMapping | None:
        z = input_data["z"]
        y_1 = input_data["y_1"]["dummy"]
        return {"y_2": array([[abs(y_1[0]) + z[0] + z[1]]])}

Define the data converters for the custom types of the variables y_1 and y_2.

This one shall return the value of a variable as a 1D array.

to_array = {
    "y_1": operator.itemgetter("dummy"),
    "y_2": operator.itemgetter(0),
}


# This one shall return the value of a variable in the type expected by the discipline from a 1D array.
def get_dict(array_):
    return {"dummy": array_}


def get_2d_array(array_):
    return array([array_])


from_array = {
    "y_1": get_dict,
    "y_2": get_2d_array,
}


# This one shall return the size of a value of a variable.
# The builtin converters already handle NumPy arrays via the ``size`` attribute,
# so we only need to handle ``y_1``.
def get_size(data):
    return data["dummy"].size


to_size = {
    "y_1": get_size,
}

set_data_converters(to_array, from_array, to_size)

disciplines = [Sellar1(), Sellar2(), SellarSystem()]

design_space = DesignSpace()
design_space.add_variable("x", lower_bound=0.0, upper_bound=10.0, value=ones(1))
design_space.add_variable(
    "z", 2, lower_bound=(-10, 0.0), upper_bound=(10.0, 10.0), value=array([4.0, 3.0])
)

scenario = create_scenario(
    disciplines,
    "obj",
    design_space,
    formulation_name="MDF",
)

scenario.add_constraint("c_1", constraint_type="ineq")
scenario.add_constraint("c_2", constraint_type="ineq")

scenario.execute(
    algo_name="NLOPT_COBYLA",
    max_iter=10,
)

# Reset the data converters to remove the custom converters.
# This should only be needed if the type of the variables y_1 and y_2 are no longer
# custom type afterward.
set_data_converters({}, {}, {})

INFO - 16:22:30: *** Start MDOScenario execution ***
INFO - 16:22:30: MDOScenario
INFO - 16:22:30:    Disciplines: Sellar1 Sellar2 SellarSystem
INFO - 16:22:30:    MDO formulation: MDF
INFO - 16:22:30: Optimization problem:
INFO - 16:22:30:    minimize obj(x, z)
INFO - 16:22:30:    with respect to x, z
INFO - 16:22:30:    under the inequality constraints
INFO - 16:22:30:       c_1(x, z) <= 0
INFO - 16:22:30:       c_2(x, z) <= 0
INFO - 16:22:30:    over the design space:
INFO - 16:22:30:       +------+-------------+-------+-------------+-------+
INFO - 16:22:30:       | Name | Lower bound | Value | Upper bound | Type  |
INFO - 16:22:30:       +------+-------------+-------+-------------+-------+
INFO - 16:22:30:       | x    |      0      |   1   |      10     | float |
INFO - 16:22:30:       | z[0] |     -10     |   4   |      10     | float |
INFO - 16:22:30:       | z[1] |      0      |   3   |      10     | float |
INFO - 16:22:30:       +------+-------------+-------+-------------+-------+
INFO - 16:22:30: Solving optimization problem with algorithm NLOPT_COBYLA:
INFO - 16:22:30:     10%|█         | 1/10 [00:00<00:00, 174.62 it/sec, feas=True, obj=21.8]
INFO - 16:22:30:     20%|██        | 2/10 [00:00<00:00, 183.44 it/sec, feas=True, obj=35.5]
INFO - 16:22:30:     30%|███       | 3/10 [00:00<00:00, 190.24 it/sec, feas=True, obj=84.8]
INFO - 16:22:30:     40%|████      | 4/10 [00:00<00:00, 192.42 it/sec, feas=True, obj=26.2]
INFO - 16:22:30:     50%|█████     | 5/10 [00:00<00:00, 191.42 it/sec, feas=True, obj=15.5]
INFO - 16:22:30:     60%|██████    | 6/10 [00:00<00:00, 192.67 it/sec, feas=True, obj=27.9]
INFO - 16:22:30:     70%|███████   | 7/10 [00:00<00:00, 192.78 it/sec, feas=True, obj=7.44]
INFO - 16:22:30:     80%|████████  | 8/10 [00:00<00:00, 189.03 it/sec, feas=True, obj=3.51]
INFO - 16:22:30:     90%|█████████ | 9/10 [00:00<00:00, 186.49 it/sec, feas=False, obj=2.81]
INFO - 16:22:30:    100%|██████████| 10/10 [00:00<00:00, 183.91 it/sec, feas=False, obj=3.68]
INFO - 16:22:30: Optimization result:
INFO - 16:22:30:    Optimizer info:
INFO - 16:22:30:       Status: None
INFO - 16:22:30:       Message: Maximum number of iterations reached. GEMSEO stopped the driver.
INFO - 16:22:30:    Solution:
INFO - 16:22:30:       The solution is feasible.
INFO - 16:22:30:       Objective: 3.506317890821701
INFO - 16:22:30:       Standardized constraints:
INFO - 16:22:30:          c_1 = -0.3267639063985981
INFO - 16:22:30:          c_2 = -20.065423793491785
INFO - 16:22:30:       Design space:
INFO - 16:22:30:          +------+-------------+-----------------------+-------------+-------+
INFO - 16:22:30:          | Name | Lower bound |         Value         | Upper bound | Type  |
INFO - 16:22:30:          +------+-------------+-----------------------+-------------+-------+
INFO - 16:22:30:          | x    |      0      |           0           |      10     | float |
INFO - 16:22:30:          | z[0] |     -10     |   2.067287914435337   |      10     | float |
INFO - 16:22:30:          | z[1] |      0      | 2.081668171172169e-16 |      10     | float |
INFO - 16:22:30:          +------+-------------+-----------------------+-------------+-------+
INFO - 16:22:30: *** End MDOScenario execution ***