Merge branch 'develop' into integration

sostrades_optimization_plugins/models/autodifferentiated_discipline.py

@@ -13,44 +13,97 @@
 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)
         self.model: Union[DifferentiableModel, None] = None
 
     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)

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.
 

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: