diff --git a/sostrades_optimization_plugins/models/autodifferentiated_discipline.py b/sostrades_optimization_plugins/models/autodifferentiated_discipline.py index 48bc68e..a2d505d 100644 --- a/sostrades_optimization_plugins/models/autodifferentiated_discipline.py +++ b/sostrades_optimization_plugins/models/autodifferentiated_discipline.py @@ -13,21 +13,22 @@ See the License for the specific language governing permissions and limitations under the License. ''' +import numpy as np import logging from typing import TYPE_CHECKING, Union from sostrades_core.execution_engine.sos_wrapp import SoSWrapp if TYPE_CHECKING: - from sostrades_optimization_plugins.models.differentiable_model import ( - DifferentiableModel, - ) + from sostrades_optimization_plugins.models.differentiable_model import DifferentiableModel class AutodifferentiedDisc(SoSWrapp): """Discipline which model is a DifferentiableModel""" + GRADIENTS = "gradient" coupling_inputs = [] # inputs verified during jacobian test coupling_outputs = [] # outputs verified during jacobian test + autoconfigure_gradient_variables: bool = True def __init__(self, sos_name, logger: logging.Logger): super().__init__(sos_name, logger) @@ -35,22 +36,74 @@ def __init__(self, sos_name, logger: logging.Logger): def run(self): - inputs = self.get_sosdisc_inputs() - self.model.set_inputs(inputs) - outputs = self.model.compute() + # todo : remove filtration later when we will be able to collect only non-numerical inputs + inputs = self.get_non_numerical_inputs() + inputs_filtered = {key: value for key, value in inputs.items() if value is not None} + self.model.set_inputs(inputs_filtered) + self.model.compute() + outputs = self.model.get_all_variables() self.store_sos_outputs_values(outputs) + def get_non_numerical_inputs(self): + inputs = self.get_sosdisc_inputs() + return {key: value for key, value in inputs.items() if value is not None} + + def compute_sos_jacobian(self): """ Compute jacobian for each coupling variable """ - gradients = self.model.compute_jacobians_custom(outputs=self.coupling_outputs, inputs=self.coupling_inputs) - for output_name in gradients: - for output_col in gradients[output_name]: - for input_name in gradients[output_name][output_col]: - for input_col, value in gradients[output_name][output_col][input_name].items(): - self.set_partial_derivative_for_other_types( - (output_name, output_col), - (input_name, input_col), - value) + if self.autoconfigure_gradient_variables: + self._auto_configure_jacobian_variables() + # dataframes variables + all_inputs_dict = {**self.DESC_IN, **self.inst_desc_in} + all_outputs_dict = {**self.DESC_OUT, **self.inst_desc_out} + coupling_dataframe_input = list(filter(lambda x: all_inputs_dict[x]['type'] == 'dataframe', self.coupling_inputs)) + coupling_dataframe_output = list(filter(lambda x: all_outputs_dict[x]['type'] == 'dataframe', self.coupling_outputs)) + other_coupling_inputs = list(set(self.coupling_inputs) - set(coupling_dataframe_input)) + other_coupling_outputs = list(set(self.coupling_outputs) - set(coupling_dataframe_output)) + + all_inputs_model_path = other_coupling_inputs + for c_i_df in coupling_dataframe_input: + all_inputs_model_path.extend(self.model.get_df_input_dotpaths(df_inputname=c_i_df)) + + all_inputs_model_path = list(filter(lambda x: not x.endswith(f":years"), all_inputs_model_path)) + + all_outputs_model_path = other_coupling_outputs + for c_o_df in coupling_dataframe_output: + all_outputs_model_path.extend(self.model.get_df_output_dotpaths(df_outputname=c_o_df)) + all_outputs_model_path = list(filter(lambda x: not x.endswith(f":years"), all_outputs_model_path)) + + + def handle_gradients_wrt_inputs(output_path: str, gradients: dict): + arg_output = (output_path,) + if ':' in output_path: + arg_output = tuple(output_path.split(':')) + + for input_path, grad_input_value in gradients.items(): + arg_input = (input_path,) + if ':' in input_path: + arg_input = tuple(input_path.split(':')) + if len(grad_input_value.shape) == 0: + grad_input_value = np.array([[grad_input_value]]) + self.set_partial_derivative_for_other_types(arg_output, arg_input, grad_input_value) + + + for output_path in all_outputs_model_path: + gradients = self.model.compute_partial(output_name=output_path, input_names=all_inputs_model_path) + handle_gradients_wrt_inputs(output_path=output_path, gradients=gradients) + + + def _auto_configure_jacobian_variables(self): + self.coupling_inputs = [] + all_inputs_dict = {**self.DESC_IN, **self.inst_desc_in} + for varname, vardescr in all_inputs_dict.items(): + if self.GRADIENTS in vardescr and vardescr[self.GRADIENTS]: + self.coupling_inputs.append(varname) + + self.coupling_outputs = [] + all_outputs_dict = {**self.DESC_OUT, **self.inst_desc_out} + for varname, vardescr in all_outputs_dict.items(): + if self.GRADIENTS in vardescr and vardescr[self.GRADIENTS]: + self.coupling_outputs.append(varname) diff --git a/sostrades_optimization_plugins/models/differentiable_model.py b/sostrades_optimization_plugins/models/differentiable_model.py index fcbe71c..558a79c 100644 --- a/sostrades_optimization_plugins/models/differentiable_model.py +++ b/sostrades_optimization_plugins/models/differentiable_model.py @@ -17,7 +17,7 @@ from collections import defaultdict from copy import deepcopy -from typing import Callable, Union +from typing import Callable, Union, Any try: import jax @@ -298,6 +298,23 @@ def get_dataframes(self, get_from: str = "outputs") -> dict[str, pd.DataFrame]: return result + def get_all_variables(self, get_from: str = "outputs") -> dict[str, Any]: + """Retrieves all variables (in or out) while converting dataframes""" + result = self.get_dataframes(get_from=get_from) + self.dataframes_outputs_colnames = self.get_output_df_names() + if get_from == "inputs": + source = self.inputs + elif get_from == "outputs": + source = self.outputs + else: + source = self.outputs + + for key, value in source.items(): + if ":" not in key and not isinstance(value, dict): + result[key] = value + + return result + def compute(self, *args: InputType) -> OutputType: """Compute the model outputs based on inputs passed as arguments.""" self._compute(*args) @@ -790,56 +807,6 @@ def check_partial( "within_tolerance": within_tolerance, } - def compute_jacobians_custom( - self, outputs: list[str], inputs: list[str] - ) -> dict[str : dict[str : dict[str : dict[str : np.ndarray]]]]: - """ - Return a dictionnary 'gradients' containing gradients for SoSwrapp disciplines. - - gradients[output df name][output column name][input df name][input column name] = value. - - """ - - # Make sure output column names are known: - self.dataframes_outputs_colnames = self.get_output_df_names() - - gradients = {} - all_inputs_paths = [] - for input_df_name in inputs: - all_inputs_paths.extend(self.get_df_input_dotpaths(input_df_name)) - all_inputs_paths = list( - filter( - lambda x: not (str(x).endswith(f":{GlossaryCore.Years}")), - all_inputs_paths, - ) - ) - for output in outputs: - gradients[output] = {} - output_columns_paths = list( - filter( - lambda x: not (str(x).endswith(f":{GlossaryCore.Years}")), - self.get_df_output_dotpaths(output), - ) - ) - for output_path in output_columns_paths: - gradients_output_path = self.compute_partial( - output_name=output_path, input_names=all_inputs_paths - ) - output_colname = output_path.split(f"{output}:")[1] - gradients[output][output_colname] = {} - for ip, value_grad in gradients_output_path.items(): - input_varname, input_varname_colname = ip.split(":") - if input_varname in gradients[output][output_colname]: - gradients[output][output_colname][input_varname][ - input_varname_colname - ] = value_grad - else: - gradients[output][output_colname][input_varname] = { - input_varname_colname: value_grad - } - - return gradients - def get_df_input_dotpaths(self, df_inputname: str) -> dict[str : list[str]]: """Get dataframe inputs dotpaths. diff --git a/sostrades_optimization_plugins/tools/discipline_tester.py b/sostrades_optimization_plugins/tools/discipline_tester.py index 2a45001..1d6b740 100644 --- a/sostrades_optimization_plugins/tools/discipline_tester.py +++ b/sostrades_optimization_plugins/tools/discipline_tester.py @@ -20,6 +20,7 @@ from sostrades_core.tests.core.abstract_jacobian_unit_test import ( AbstractJacobianUnittest, ) +from sostrades_optimization_plugins.models.autodifferentiated_discipline import AutodifferentiedDisc def discipline_test_function(module_path: str, name: str, model_name: str, @@ -64,6 +65,12 @@ def discipline_test_function(module_path: str, name: str, model_name: str, filter = disc.get_chart_filter_list() graph_list = disc.get_post_processing_list(filter) + wrap_disc = disc.discipline_wrapp.wrapper + if not coupling_inputs and not coupling_outputs: + if isinstance(wrap_disc, AutodifferentiedDisc) and wrap_disc.autoconfigure_gradient_variables: + wrap_disc._auto_configure_jacobian_variables() + coupling_inputs, coupling_outputs = wrap_disc.coupling_inputs, wrap_disc.coupling_outputs + # Show generated graphs if show_graphs: for graph in graph_list: