Merge pull request #373 from Ipuch/solParameters

Displaying optimized parameters in sol.print()

Author: pariterre

bioptim/dynamics/configure_problem.py

@@ -502,7 +502,7 @@ def configure_tau(nlp, as_states: bool, as_controls: bool, fatigue: FatigueList
                 )
 
                 for p, params in enumerate(fatigue_suffix):
-                    name = f"tau_{params}_{tau_suffix}"
+                    name = f"tau_{tau_suffix}_{params}"
                     ConfigureProblem._adjust_mapping(name, ["q"], nlp)
                     ConfigureProblem.configure_new_variable(name, name_tau, nlp, True, False, skip_plot=True)
                     nlp.plot[f"{name}_controls"] = CustomPlot(

bioptim/dynamics/dynamics_functions.py

@@ -159,7 +159,7 @@ def __get_fatigable_tau(nlp: NonLinearProgram, states: MX, controls: MX, fatigue
                     for suffix in tau_suffix:
                         model = getattr(t.model, suffix)
                         tau_tp += (
-                            DynamicsFunctions.get(nlp.states[f"tau_{model.dynamics_suffix()}_{suffix}"], states)[i]
+                            DynamicsFunctions.get(nlp.states[f"tau_{suffix}_{model.dynamics_suffix()}"], states)[i]
                             * model.scale
                         )
                     tau = vertcat(tau, tau_tp)

bioptim/dynamics/xia_fatigue.py

@@ -121,13 +121,13 @@ def _dynamics_per_suffix(self, dxdt, suffix, nlp, index, states, controls):
         var = getattr(self, suffix)
         target_load = self._get_target_load(var, suffix, nlp, controls, index)
         fatigue = [
-            DynamicsFunctions.get(nlp.states[f"tau_{dyn_suffix}_{suffix}"], states)[index, :]
+            DynamicsFunctions.get(nlp.states[f"tau_{suffix}_{dyn_suffix}"], states)[index, :]
             for dyn_suffix in var.suffix()
         ]
         current_dxdt = var.apply_dynamics(target_load, *fatigue)
 
         for i, dyn_suffix in enumerate(var.suffix()):
-            dxdt[nlp.states[f"tau_{dyn_suffix}_{suffix}"].index[index], :] = current_dxdt[i]
+            dxdt[nlp.states[f"tau_{suffix}_{dyn_suffix}"].index[index], :] = current_dxdt[i]
 
         return dxdt
 

bioptim/gui/plot.py

@@ -58,7 +58,7 @@ def __init__(
         bounds: Bounds = None,
         node_idx: list = None,
         label: list = None,
-        manually_compute_derivative: bool = False,
+        compute_derivative: bool = False,
         **parameters: Any,
     ):
         """
@@ -86,7 +86,7 @@ def __init__(
             The node time to be plotted on the graphs
         label: list
             Label of the curve to plot (to be added to the legend)
-        manually_compute_derivative: bool
+        compute_derivative: bool
             If the function should send the next node with x and u. Prevents from computing all at once (therefore a bit slower)
         """
 
@@ -108,7 +108,7 @@ def __init__(
         self.bounds = bounds
         self.node_idx = node_idx
         self.label = label
-        self.manually_compute_derivative = manually_compute_derivative
+        self.compute_derivative = compute_derivative
         self.parameters = parameters
 
 
@@ -614,21 +614,14 @@ def update_data(self, v: dict):
                 else:
                     control = np.concatenate((control, data_controls[s]))
 
-            if nlp.control_type == ControlType.CONSTANT:
-                u_mod = 1
-            elif nlp.control_type == ControlType.LINEAR_CONTINUOUS:
-                u_mod = 2
-            else:
-                raise NotImplementedError(f"Plotting {nlp.control_type} is not implemented yet")
-
             for key in self.variable_sizes[i]:
                 if not self.plot_func[key][i]:
                     continue
-                # Automatically find u_modifier if the function is a casadi function otherwise fallback to default
-                u_mod2 = (
-                    self.plot_func[key][i].function.size2_in(1)
-                    if hasattr(self.plot_func[key][i].function, "size2_in")
-                    else u_mod
+                x_mod = 1 if self.plot_func[key][i].compute_derivative else 0
+                u_mod = (
+                    1
+                    if nlp.control_type == ControlType.LINEAR_CONTINUOUS or self.plot_func[key][i].compute_derivative
+                    else 0
                 )
 
                 if self.plot_func[key][i].type == PlotType.INTEGRATED:
@@ -637,17 +630,21 @@ def update_data(self, v: dict):
                         y_tp = np.empty((self.variable_sizes[i][key], len(t)))
                         y_tp.fill(np.nan)
 
-                        mod = 1 if self.plot_func[key][i].manually_compute_derivative else 0
                         val = self.plot_func[key][i].function(
                             idx,
-                            state[:, step_size * idx : step_size * (idx + 1) + mod],
-                            control[:, idx : idx + u_mod2 + 1],
+                            state[:, step_size * idx : step_size * (idx + 1) + x_mod],
+                            control[:, idx : idx + u_mod + 1],
                             data_params_in_dyn,
                             **self.plot_func[key][i].parameters,
                         )
+
+                        if self.plot_func[key][i].compute_derivative:
+                            # This is a special case since derivative is not properly integrated
+                            val = np.repeat(val, y_tp.shape[1])[np.newaxis, :]
+
                         if val.shape != y_tp.shape:
                             raise RuntimeError(
-                                f"Wrong dimensions for plot {key}. Got {val.shape}, but expected {y.shape}"
+                                f"Wrong dimensions for plot {key}. Got {val.shape}, but expected {y_tp.shape}"
                             )
                         y_tp[:, :] = val
                         all_y.append(y_tp)
@@ -659,24 +656,25 @@ def update_data(self, v: dict):
                         self.__append_to_ydata([y_tp])
 
                 elif self.plot_func[key][i].type == PlotType.POINT:
-                    y = np.empty((len(self.plot_func[key][i].node_idx),))
-                    y.fill(np.nan)
-                    mod = 1 if self.plot_func[key][i].manually_compute_derivative else 0
-                    for i_node, node_idx in enumerate(self.plot_func[key][i].node_idx):
-                        val = self.plot_func[key][i].function(
-                            node_idx,
-                            state[:, node_idx * step_size : (node_idx + 1) * step_size + mod : step_size],
-                            control[:, node_idx : node_idx + 1 + mod],
-                            data_params_in_dyn,
-                            **self.plot_func[key][i].parameters,
-                        )
-                        y[i_node] = val
-                    self.ydata.append(y)
+                    for i_var in range(self.variable_sizes[i][key]):
+                        y = np.empty((len(self.plot_func[key][i].node_idx),))
+                        y.fill(np.nan)
+                        mod = 1 if self.plot_func[key][i].compute_derivative else 0
+                        for i_node, node_idx in enumerate(self.plot_func[key][i].node_idx):
+                            val = self.plot_func[key][i].function(
+                                node_idx,
+                                state[:, node_idx * step_size : (node_idx + 1) * step_size + mod : step_size],
+                                control[:, node_idx : node_idx + 1 + mod],
+                                data_params_in_dyn,
+                                **self.plot_func[key][i].parameters,
+                            )
+                            y[i_node] = val[i_var]
+                        self.ydata.append(y)
 
                 else:
                     y = np.empty((self.variable_sizes[i][key], len(self.t[i])))
                     y.fill(np.nan)
-                    if self.plot_func[key][i].manually_compute_derivative:
+                    if self.plot_func[key][i].compute_derivative:
                         for i_node, node_idx in enumerate(self.plot_func[key][i].node_idx):
                             val = self.plot_func[key][i].function(
                                 node_idx,
@@ -687,8 +685,18 @@ def update_data(self, v: dict):
                             )
                             y[:, i_node] = val
                     else:
+                        nodes = self.plot_func[key][i].node_idx
+                        if nodes and len(nodes) > 1 and len(nodes) == round(state.shape[1] / step_size):
+                            # Assume we are integrating but did not specify plot as such.
+                            # Therefore the arrival point is missing
+                            nodes += [nodes[-1] + 1]
+
                         val = self.plot_func[key][i].function(
-                            i, state[:, ::step_size], control, data_params_in_dyn, **self.plot_func[key][i].parameters
+                            nodes,
+                            state[:, ::step_size],
+                            control,
+                            data_params_in_dyn,
+                            **self.plot_func[key][i].parameters,
                         )
                         if val.shape != y.shape:
                             raise RuntimeError(

bioptim/limits/penalty_option.py

@@ -161,6 +161,7 @@ def __init__(
         self.weight = weight
         self.function: Union[Function, None] = None
         self.weighted_function: Union[Function, None] = None
+        self.weighted_function_non_threaded: Union[Function, None] = None
         self.derivative = derivative
         self.explicit_derivative = explicit_derivative
         self.integrate = integrate
@@ -337,6 +338,7 @@ def _set_penalty_function(self, all_pn: Union[PenaltyNodeList, list, tuple], fcn
         self.weighted_function = Function(
             name, [state_cx, control_cx, param_cx, weight_cx, target_cx, dt_cx], [modified_fcn]
         )
+        self.weighted_function_non_threaded = self.weighted_function
 
         if ocp.n_threads > 1 and self.multi_thread and len(self.node_idx) > 1:
             self.function = self.function.map(len(self.node_idx), "thread", ocp.n_threads)
@@ -380,14 +382,21 @@ def _finish_add_target_to_plot(self, all_pn: PenaltyNodeList):
 
         """
 
+        def plot_function(t, x, u, p):
+            if isinstance(t, (list, tuple)):
+                return self.target_to_plot[:, [self.node_idx.index(_t) for _t in t]]
+            else:
+                return self.target_to_plot[:, self.node_idx.index(t)]
+
         if self.target_to_plot is not None:
-            if self.target_to_plot.shape[0] > 1:
+            if self.target_to_plot.shape[1] > 1:
                 plot_type = PlotType.STEP
             else:
                 plot_type = PlotType.POINT
+
             all_pn.ocp.add_plot(
                 self.target_plot_name,
-                lambda t, x, u, p: self.target_to_plot,
+                plot_function,
                 color="tab:red",
                 plot_type=plot_type,
                 phase=all_pn.nlp.phase_idx,

bioptim/optimization/optimal_control_program.py

@@ -647,11 +647,11 @@ def penalty_color():
                     penalties_internal = nlp.g_internal
 
                 for penalty in penalties:
-                    if penalty is None:
+                    if not penalty:
                         continue
                     name_unique_objective.append(penalty.name)
                 for penalty_internal in penalties_internal:
-                    if penalty_internal is None:
+                    if not penalty_internal:
                         continue
                     name_unique_objective.append(penalty_internal.name)
             color = {}
@@ -671,21 +671,24 @@ def compute_penalty_values(t, x, u, p, penalty, dt):
                 if dt.shape[0] > 1:
                     dt = dt[penalty.phase]
 
-            _target = penalty.target[..., t] if penalty.target is not None and isinstance(t, int) else []
+            _target = (
+                penalty.target[..., penalty.node_idx.index(t)]
+                if penalty.target is not None and isinstance(t, int)
+                else []
+            )
 
             out = []
             if penalty.transition:
                 raise NotImplementedError("add_plot_penalty with phase transition is not implemented yet")
             elif penalty.derivative or penalty.explicit_derivative:
-                out.append(penalty.weighted_function(x, u, p, penalty.weight, _target, dt))
+                out.append(penalty.weighted_function_non_threaded(x[:, [0, -1]], u, p, penalty.weight, _target, dt))
             else:
-                _u = u if penalty.weighted_function.sparsity_in(1).shape[1] > 1 else u[:, :-1]
-                out.append(penalty.weighted_function(x[:, :-1], _u, p, penalty.weight, _target, dt))
+                out.append(penalty.weighted_function_non_threaded(x, u, p, penalty.weight, _target, dt))
             return sum1(horzcat(*out))
 
         def add_penalty(_penalties):
             for penalty in _penalties:
-                if penalty is None:
+                if not penalty:
                     continue
 
                 dt = penalty.dt
@@ -707,7 +710,7 @@ def add_penalty(_penalties):
                     "dt": dt,
                     "color": color[penalty.name],
                     "label": penalty.name,
-                    "manually_compute_derivative": True,
+                    "compute_derivative": penalty.derivative or penalty.explicit_derivative or penalty.integrate,
                 }
                 if (
                     isinstance(penalty.type, ObjectiveFcn.Mayer)

tests/test_global_fatigue.py

@@ -102,8 +102,8 @@ def test_fatigable_torque():
     # Check some of the results
     states, controls = sol.states, sol.controls
     q, qdot = states["q"], states["qdot"]
-    ma_minus, mr_minus, mf_minus = states["tau_ma_minus"], states["tau_mr_minus"], states["tau_mf_minus"]
-    ma_plus, mr_plus, mf_plus = states["tau_ma_plus"], states["tau_mr_plus"], states["tau_mf_plus"]
+    ma_minus, mr_minus, mf_minus = states["tau_minus_ma"], states["tau_minus_mr"], states["tau_minus_mf"]
+    ma_plus, mr_plus, mf_plus = states["tau_plus_ma"], states["tau_plus_mr"], states["tau_plus_mf"]
     tau_minus, tau_plus = controls["tau_minus"], controls["tau_plus"]
 
     # initial and final position