Merge pull request #1048 from pariterre/master

Formalized parameters optimization (#1048)

Author: EveCharbie

.vscode/settings.json

@@ -10,5 +10,7 @@
         "editor.rulers": [
             120
         ]
-    }
+    },
+    "python-envs.defaultEnvManager": "ms-python.python:conda",
+    "python-envs.defaultPackageManager": "ms-python.python:conda"
 }

bioptim/__init__.py

@@ -239,7 +239,7 @@
 from .optimization.optimal_control_program import OptimalControlProgram
 from .optimization.optimization_variable import OptimizationVariableList
 from .optimization.vector_layout import OrderingStrategy
-from .optimization.parameters import ParameterList, ParameterContainer
+from .optimization.parameters import Parameter, ParameterList, ParameterContainer
 from .optimization.problem_type import SocpType
 from .optimization.receding_horizon_optimization import (
     CyclicNonlinearModelPredictiveControl,

bioptim/dynamics/dynamics_functions.py

@@ -1,3 +1,5 @@
+from typing import TYPE_CHECKING
+
 from casadi import vertcat, MX, SX
 
 from ..models.protocols.holonomic_constraints import HolonomicConstraintsFcn
@@ -6,6 +8,9 @@
 from ..optimization.optimization_variable import OptimizationVariable
 from ..misc.parameters_types import AnyListOptional, CX, CXOptional, Str, Tuple
 
+if TYPE_CHECKING:
+    from ..optimization.non_linear_program import NonLinearProgram
+
 
 class DynamicsFunctions:
     """
@@ -26,7 +31,7 @@ class DynamicsFunctions:
     """
 
     @staticmethod
-    def get_fatigable_tau(nlp, states: CX, controls: CX) -> CX:
+    def get_fatigable_tau(nlp: "NonLinearProgram", states: CX, controls: CX) -> CX:
         """
         Apply the forward dynamics including (or not) the torque fatigue
 
@@ -82,7 +87,7 @@ def get_fatigable_tau(nlp, states: CX, controls: CX) -> CX:
     @staticmethod
     def get_fatigue_states(
         states,
-        nlp,
+        nlp: "NonLinearProgram",
         mus_activations,
     ):
 
@@ -141,7 +146,7 @@ def get(var: OptimizationVariable, cx: CX):
         return var.mapping.to_second.map(cx[var.index, :])
 
     @staticmethod
-    def compute_qdot(nlp, q: CX, qdot: CX):
+    def compute_qdot(nlp: "NonLinearProgram", q: CX, qdot: CX):
         """
         Easy accessor to derivative of q
 
@@ -175,7 +180,7 @@ def compute_qdot(nlp, q: CX, qdot: CX):
         return mapping.to_first.map(nlp.model.reshape_qdot()(q, qdot, nlp.parameters.cx))
 
     @staticmethod
-    def compute_qddot(nlp, q: CX, qdot: CX, tau: CX, external_forces: CX):
+    def compute_qddot(nlp: "NonLinearProgram", q: CX, qdot: CX, tau: CX, external_forces: CX):
         """
         Easy accessor to derivative of qdot
 
@@ -202,7 +207,7 @@ def compute_qddot(nlp, q: CX, qdot: CX, tau: CX, external_forces: CX):
 
     @staticmethod
     def collect_tau(
-        nlp,
+        nlp: "NonLinearProgram",
         q: CX,
         qdot: CX,
         parameters: CX,
@@ -231,7 +236,7 @@ def collect_tau(
         return tau
 
     @staticmethod
-    def get_contact_defects(nlp, q: CX, qdot: CX, slope_qdot: CX):
+    def get_contact_defects(nlp: "NonLinearProgram", q: CX, qdot: CX, slope_qdot: CX):
         """
         Get the defects associated with implicit contacts.
 
@@ -263,7 +268,9 @@ def get_contact_defects(nlp, q: CX, qdot: CX, slope_qdot: CX):
         return contact_defects
 
     @staticmethod
-    def get_fatigue_defects(key: Str, dxdt_defects: CX, slopes: CX, nlp, states: CX, controls: CX) -> Tuple[CX, CX]:
+    def get_fatigue_defects(
+        key: Str, dxdt_defects: CX, slopes: CX, nlp: "NonLinearProgram", states: CX, controls: CX
+    ) -> Tuple[CX, CX]:
         """
         Get the dxdt and slopes associated with fatigue elements.
         These are added to compute the defects in the case where there is fatigue.
@@ -302,7 +309,7 @@ def get_fatigue_defects(key: Str, dxdt_defects: CX, slopes: CX, nlp, states: CX,
         return dxdt_defects, slopes
 
     @staticmethod
-    def get_external_forces_from_contacts(nlp, q, qdot, contact_types, external_forces: MX | SX):
+    def get_external_forces_from_contacts(nlp: "NonLinearProgram", q, qdot, contact_types, external_forces: MX | SX):
         """
         Get the external forces associated with the contacts defined in the model.
 
@@ -353,7 +360,7 @@ def get_external_forces_from_contacts(nlp, q, qdot, contact_types, external_forc
 
     @staticmethod
     def forward_dynamics(
-        nlp,
+        nlp: "NonLinearProgram",
         q: CX,
         qdot: CX,
         tau: CX,
@@ -412,7 +419,7 @@ def forward_dynamics(
 
     @staticmethod
     def inverse_dynamics(
-        nlp,
+        nlp: "NonLinearProgram",
         q: CX,
         qdot: CX,
         qddot: CX,
@@ -472,7 +479,7 @@ def inverse_dynamics(
         return tau_var_mapping.map(tau)
 
     @staticmethod
-    def compute_muscle_dot(nlp, muscle_excitations: CX, muscle_activations: CX):
+    def compute_muscle_dot(nlp: "NonLinearProgram", muscle_excitations: CX, muscle_activations: CX):
         """
         Easy accessor to derivative of muscle activations
 
@@ -494,7 +501,7 @@ def compute_muscle_dot(nlp, muscle_excitations: CX, muscle_activations: CX):
 
     @staticmethod
     def compute_tau_from_muscle(
-        nlp,
+        nlp: "NonLinearProgram",
         q: CX,
         qdot: CX,
         muscle_activations: CX,
@@ -530,7 +537,7 @@ def compute_tau_from_muscle(
         return nlp.model.muscle_joint_torque()(activations, q, qdot, nlp.parameters.cx)
 
     @staticmethod
-    def no_states_mapping(nlp):
+    def no_states_mapping(nlp: "NonLinearProgram"):
         for key in nlp.states.keys():
             if nlp.variable_mappings[key].actually_does_a_mapping():
                 raise NotImplementedError(

bioptim/examples/getting_started/custom_parameters.py

@@ -28,11 +28,12 @@
     ObjectiveList,
     PhaseDynamics,
     VariableScaling,
+    Parameter,
 )
 from bioptim.examples.utils import ExampleUtils
 
 
-def my_parameter_function(bio_model: TorqueBiorbdModel, value: MX, extra_value: Any):
+def my_parameter_function(bio_model: TorqueBiorbdModel, parameter: Parameter, extra_value: Any):
     """
     The pre dynamics function is called right before defining the dynamics of the system. If one wants to
     modify the dynamics (e.g. optimize the gravity in this case), then this function is the proper way to do it.
@@ -41,17 +42,17 @@ def my_parameter_function(bio_model: TorqueBiorbdModel, value: MX, extra_value:
     ----------
     bio_model: TorqueBiorbdModel
         The model to modify by the parameters
-    value: MX
-        The CasADi variables to modify the model
+    parameter: Parameter
+        The optimization variables to modify the model
     extra_value: Any
         Any parameters required by the user. The name(s) of the extra_value must match those used in parameter.add
     """
 
-    value[2] *= extra_value
-    bio_model.set_gravity(value)
+    parameter.mx[2] *= extra_value
+    bio_model.set_gravity(parameter)
 
 
-def set_mass(bio_model: TorqueBiorbdModel, value: MX):
+def set_mass(bio_model: TorqueBiorbdModel, parameter: Parameter):
     """
     The pre dynamics function is called right before defining the dynamics of the system. If one wants to
     modify the dynamics (e.g. optimize the gravity in this case), then this function is the proper way to do it.
@@ -60,11 +61,11 @@ def set_mass(bio_model: TorqueBiorbdModel, value: MX):
     ----------
     bio_model: TorqueBiorbdModel
         The model to modify by the parameters
-    value: MX
-        The CasADi variables to modify the model
+    parameter: Parameter
+        The optimization variables to modify the model
     """
-
-    bio_model.segments[0].characteristics().setMass(value)
+    # Since Biorbd uses MX internally, we need to use the MX version of the variable, which is accessible through the .mx attribute
+    bio_model.segments[0].characteristics().setMass(parameter.mx)
 
 
 def my_target_function(controller: PenaltyController, key: str) -> MX:

bioptim/examples/toy_examples/feature_examples/example_parameter_scaling.py

@@ -25,6 +25,7 @@
     PhaseDynamics,
     VariableScaling,
     SolutionMerge,
+    Parameter,
 )
 from bioptim.examples.utils import ExampleUtils
 import numpy as np
@@ -106,8 +107,8 @@ def generate_dat_to_track(
     )
 
 
-def my_parameter_function(bio_model: TorqueBiorbdModel, value: MX):
-    bio_model.set_gravity(value)
+def my_parameter_function(bio_model: TorqueBiorbdModel, parameter: Parameter):
+    bio_model.set_gravity(parameter)
 
 
 def my_target_function(controller: PenaltyController, key: str) -> MX:

bioptim/examples/toy_examples/holonomic_constraints/custom_dynamics.py

@@ -252,19 +252,18 @@ class HolonomicMusclesBiorbdModel(HolonomicBiorbdModel, HolonomicMusclesDynamics
     def __init__(
         self,
         bio_model: str | biorbd.Model,
-        friction_coefficients: np.ndarray = None,
         parameters: ParameterList = None,
         holonomic_constraints: HolonomicConstraintsList | None = None,
         dependent_joint_index: list[int] | tuple[int, ...] = None,
         independent_joint_index: list[int] | tuple[int, ...] = None,
     ):
-        HolonomicBiorbdModel.__init__(self, bio_model, friction_coefficients, parameters)
+        HolonomicBiorbdModel.__init__(self, bio_model, parameters)
         if holonomic_constraints is not None:
             self.set_holonomic_configuration(holonomic_constraints, dependent_joint_index, independent_joint_index)
         HolonomicMusclesDynamics.__init__(self)
 
     def serialize(self) -> tuple[Callable, dict]:
-        return HolonomicMusclesBiorbdModel, dict(bio_model=self.path, friction_coefficients=self.friction_coefficients)
+        return HolonomicMusclesBiorbdModel, dict(bio_model=self.path)
 
 
 class AlgebraicHolonomicMusclesBiorbdModel(HolonomicMusclesBiorbdModel):

bioptim/examples/toy_examples/moving_horizon_estimation/multi_cyclic_nmpc_with_parameters.py

@@ -33,6 +33,7 @@
     Solver,
     VariableScaling,
     OnlineOptim,
+    Parameter,
 )
 from bioptim.examples.utils import ExampleUtils
 from casadi import MX, SX, vertcat
@@ -55,7 +56,7 @@ def advance_window_initial_guess_states(self, sol, n_cycles_simultaneous=None):
         return True
 
 
-def dummy_parameter_function(bio_model, value: MX):
+def dummy_parameter_function(bio_model, parameter: Parameter):
     return
 
 

bioptim/examples/toy_examples/stochastic_optimal_control/arm_reaching_torque_driven_collocations.py

@@ -25,6 +25,9 @@
     Axis,
     SolutionMerge,
     DynamicsOptions,
+    Parameter,
+    ParameterList,
+    ParameterObjectiveList,
 )
 from bioptim.examples.toy_examples.stochastic_optimal_control.arm_reaching_torque_driven_implicit import ExampleType
 from bioptim.examples.utils import ExampleUtils
@@ -52,6 +55,10 @@ def sensory_reference(
     return hand_pos_velo
 
 
+def set_friction_coefficients(bio_model, parameter: Parameter):
+    bio_model.set_friction_coefficients(parameter)
+
+
 def prepare_socp(
     biorbd_model_path: str,
     final_time: float,
@@ -64,6 +71,7 @@ def prepare_socp(
     q_opt: np.ndarray = None,
     qdot_opt: np.ndarray = None,
     tau_opt: np.ndarray = None,
+    optimize_friction_coefficients: bool = False,
     use_sx=False,
 ) -> StochasticOptimalControlProgram:
     """
@@ -86,6 +94,8 @@ def prepare_socp(
         The magnitude of the sensory noise
     example_type
         The type of problem to solve (CIRCLE or BAR)
+    optimize_friction_coefficients: bool
+        If the friction coefficients should be optimized or not
     use_sx: bool
         If SX should be used instead of MX
 
@@ -103,6 +113,27 @@ def prepare_socp(
         initial_cov=initial_cov,
     )
 
+    if optimize_friction_coefficients:
+        parameters = ParameterList(use_sx=use_sx)
+        parameters.add("friction_coefficients", set_friction_coefficients, size=4)
+
+        target = np.array(2 * np.array([0.05, 0.025, 0.025, 0.05]))
+        parameter_objectives = ParameterObjectiveList()
+        parameter_objectives.add(
+            ObjectiveFcn.Parameter.MINIMIZE_PARAMETER,
+            key="friction_coefficients",
+            target=target,
+            quadratic=True,
+        )
+
+        parameter_init = InitialGuessList()
+        parameter_init.add("friction_coefficients", initial_guess=target)
+
+    else:
+        parameters = None
+        parameter_objectives = None
+        parameter_init = None
+
     bio_model = StochasticTorqueBiorbdModel(
         biorbd_model_path,
         problem_type=problem_type,
@@ -114,10 +145,14 @@ def prepare_socp(
         sensory_reference=sensory_reference,
         n_references=4,  # This number must be in agreement with what is declared in sensory_reference
         n_feedbacks=4,
+        parameters=parameters,
+        parameter_init=parameter_init,
         use_sx=use_sx,
-        friction_coefficients=np.array([[0.05, 0.025], [0.025, 0.05]]),
     )
 
+    if not optimize_friction_coefficients:
+        bio_model.set_friction_coefficients(np.array([[0.05, 0.025], [0.025, 0.05]]))
+
     n_tau = bio_model.nb_tau
     n_q = bio_model.nb_q
     n_qdot = bio_model.nb_qdot
@@ -255,10 +290,13 @@ def prepare_socp(
         x_init=x_init,
         u_init=u_init,
         a_init=a_init,
+        parameters=parameters,
         x_bounds=x_bounds,
         u_bounds=u_bounds,
         a_bounds=a_bounds,
         objective_functions=objective_functions,
+        parameter_objectives=parameter_objectives,
+        parameter_init=parameter_init,
         constraints=constraints,
         control_type=ControlType.CONSTANT_WITH_LAST_NODE,
         n_threads=1,

bioptim/examples/toy_examples/stochastic_optimal_control/arm_reaching_torque_driven_explicit.py

@@ -396,11 +396,13 @@ def prepare_socp(
         n_noised_controls=2,
         sensory_noise_magnitude=sensory_noise_magnitude,
         motor_noise_magnitude=motor_noise_magnitude,
-        friction_coefficients=np.array([[0.05, 0.025], [0.025, 0.05]]),
         sensory_reference=sensory_reference,
     )
     bio_model.force_field_magnitude = force_field_magnitude
 
+    # Please refer to the arm_reaching_torque_driven_collocations.py example for an example of optimizing these values
+    bio_model.set_friction_coefficients(np.array([[0.05, 0.025], [0.025, 0.05]]))
+
     n_tau = bio_model.nb_tau
     n_q = bio_model.nb_q
     n_qdot = bio_model.nb_qdot

bioptim/examples/toy_examples/stochastic_optimal_control/arm_reaching_torque_driven_implicit.py

@@ -98,6 +98,8 @@ def prepare_socp(
         The type of problem to solve (CIRCLE or BAR)
     with_cholesky: bool
         If True, whether to use the Cholesky factorization of the covariance matrix or not
+    with_scaling: bool
+        If True, whether to use scaling for the controls or not
 
     Returns
     -------
@@ -117,9 +119,11 @@ def prepare_socp(
         n_feedbacks=4,
         n_noised_states=4,
         n_noised_controls=2,
-        friction_coefficients=np.array([[0.05, 0.025], [0.025, 0.05]]),
     )
 
+    # Please refer to the arm_reaching_torque_driven_collocations.py example for an example of optimizing these values
+    bio_model.set_friction_coefficients(np.array([[0.05, 0.025], [0.025, 0.05]]))
+
     n_tau = bio_model.nb_tau
     n_q = bio_model.nb_q
     n_qdot = bio_model.nb_qdot