Rename the input and output variables

The RemappingDiscipline can be used to rename the input and output variables of an original discipline including defining a variable as a part of an original one.

from __future__ import annotations

from typing import TYPE_CHECKING

from numpy import array

from gemseo.core.discipline import Discipline
from gemseo.disciplines.remapping import RemappingDiscipline

if TYPE_CHECKING:
    from gemseo.typing import StrKeyMapping

Let us consider a discipline that sums up the fruits of the market. The fruits can be classified into three categories: pears, Gala apples and Fuji apples. Then, the input variable of the discipline called fruits is a triplet containing the numbers of pears, Gala apples and Fuji apples so ordered. Concerning the outputs, n_fruits is the total number of fruits while n_fruits_per_category gathers the numbers of pears and apples so ordered. This discipline can be coded as follows:

class FruitCounting(Discipline):
    def __init__(self) -> None:
        super().__init__()
        self.input_grammar.update_from_names(["fruits"])
        self.output_grammar.update_from_names(["n_fruits", "n_fruits_per_category"])
        self.default_input_data = {"fruits": array([1, 2, 3])}

    def _run(self, input_data: StrKeyMapping) -> StrKeyMapping | None:
        fruits = input_data["fruits"]
        return {
            "n_fruits": array([fruits.sum()]),
            "n_fruits_per_category": array([fruits[0], fruits[1:3].sum()]),
        }

and we can instantiate it:

fruit_counting = FruitCounting()

Then, we create a new discipline renaming fruits as pear and apples and n_fruits and n_fruits_per_category as total and sub_total to improve the naming:

clearer_fruit_counting = RemappingDiscipline(
    fruit_counting,
    {"pear": ("fruits", 0), "apples": ("fruits", [1, 2])},
    {"total": "n_fruits", "sub_total": "n_fruits_per_category"},
)

Note:

RemappingDiscipline requires an instance of the original discipline, the input names mapping to the original input names and the outputs names mapping to the original output names. More precisely, an input or output name mapping looks like {"new_x": "x", "new_y": ("y", components)} where the variable "new_x" corresponds to the original variable "x" and the variable "new_y" corresponds to some components of the original variable "y". components can be an integer i (the i-th component of y), a sequence of integers [i, j, k] (the i-th, j-th and k-th components of y) or an iterable of integers range(i, j+1) (from the i-th to the j-th components of y).

We can execute this discipline with the original default input values, namely 1 pear, 2 Gala apples and 3 Fuji apples:

clearer_fruit_counting.execute()
clearer_fruit_counting.get_input_data(), clearer_fruit_counting.get_output_data()

({'pear': array([1]), 'apples': array([2, 3])}, {'total': array([6]), 'sub_total': array([1, 5])})

or with new input data:

clearer_fruit_counting.execute({"pear": array([4]), "apples": array([3, 1])})
clearer_fruit_counting.get_input_data(), clearer_fruit_counting.get_output_data()

({'pear': array([4]), 'apples': array([3, 1])}, {'total': array([8]), 'sub_total': array([4, 4])})

To be even more clear, we can split apples into gala and fuji and sub_total into n_pears and n_apples:

even_clearer_fruit_counting = RemappingDiscipline(
    clearer_fruit_counting,
    {"pear": "pear", "gala": ("apples", 0), "fuji": ("apples", 0)},
    {
        "total": "total",
        "n_pears": ("sub_total", 0),
        "n_apples": ("sub_total", 1),
    },
)

and count the number of fruits:

even_clearer_fruit_counting.execute({
    "pear": array([4]),
    "gala": array([3]),
    "fuji": array([1]),
})
(
    even_clearer_fruit_counting.get_input_data(),
    even_clearer_fruit_counting.get_output_data(),
)

({'pear': array([4]), 'gala': array([3]), 'fuji': array([1])}, {'total': array([10]), 'n_pears': array([4]), 'n_apples': array([6])})