Source code for gemseo.disciplines.splitter

# Copyright 2021 IRT Saint Exupéry,
# 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
# 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.
"""A discipline splitting an input variable."""
from __future__ import annotations

from typing import Iterable

from numpy import eye

from gemseo.core.discipline import MDODiscipline

[docs]class Splitter(MDODiscipline): """A discipline splitting an input variable. Several output variables containing slice of the input variable are extracted. Example: >>> discipline = Splitter("alpha", {"beta": [0, 1], "delta": [2, 3], "gamma": 4}) >>> discipline.execute({"alpha": array([1.0, 2.0, 3.0, 4.0, 5.0])}) >>> delta = discipline.local_data["delta"] # delta = array([3.0, 4.0]) """ def __init__( self, input_name: str, output_names_to_input_indices: dict[str, Iterable[int] | int], ): """# noqa: D205 D212 D415 Args: input_name: The name of the input to split. output_names_to_input_indices: The input indices associated with the output names. """ self.__input_name = input_name for output_name, input_indices in output_names_to_input_indices.items(): if isinstance(input_indices, int): output_names_to_input_indices[output_name] = [input_indices] self.__slicing_structure = output_names_to_input_indices super().__init__() self.input_grammar.update([input_name]) self.output_grammar.update(output_names_to_input_indices.keys()) def _run(self) -> None: input_data = self.local_data[self.__input_name] for output_name, input_indices in self.__slicing_structure.items(): self.local_data[output_name] = input_data[input_indices] def _compute_jacobian( self, inputs: Iterable[str] | None = None, outputs: Iterable[str] | None = None, ) -> None: self._init_jacobian(with_zeros=True) self.jac = {} identity = eye(self.local_data[self.__input_name].size) for output_name, input_indices in self.__slicing_structure.items(): self.jac[output_name] = {} self.jac[output_name][self.__input_name] = identity[input_indices, :]