made requested changes

.gitignore

@@ -118,6 +118,4 @@ bioptim/sandbox/
 # Mac dev
 *.DS_store
 
-bioptim/examples/stochastic_optimal_control/output.png
-
 *.png

bioptim/dynamics/dynamics_functions.py

@@ -231,7 +231,7 @@ def stochastic_torque_driven(
 
         ref = DynamicsFunctions.get(nlp.stochastic_variables["ref"], stochastic_variables)
         k = DynamicsFunctions.get(nlp.stochastic_variables["k"], stochastic_variables)
-        k_matrix = StochasticBioModel.reshape_sym_to_matrix(k, nlp.model.matrix_shape_k)
+        k_matrix = StochasticBioModel.reshape_to_matrix(k, nlp.model.matrix_shape_k)
 
         sensory_input = nlp.model.sensory_reference(states, controls, parameters, stochastic_variables, nlp)
 

bioptim/dynamics/integrator.py

@@ -922,15 +922,6 @@ def _finish_init(self):
         """
         Prepare the CasADi function from dxdt
         """
-        xf, xall = self.dxdt(
-            h=self.h,
-            time=self.time_integration_grid[0],
-            states=self.x_sym,
-            controls=self.u_sym,
-            params=self.param_sym,
-            param_scaling=self.param_scaling,
-            stochastic_variables=self.s_sym,
-        )
 
         self.function = Function(
             "integrator",

bioptim/dynamics/ode_solver.py

@@ -314,7 +314,7 @@ class COLLOCATION(OdeSolverBase):
             The method of interpolation ("legendre" or "radau")
         defects_type: DefectType
             The type of defect to use (DefectType.EXPLICIT or DefectType.IMPLICIT)
-        add_initial_collocation_point: bool
+        include_starting_collocation_point: bool
             Whether an additional collocation point should be added at the shooting node (this is typically used in SOCPs)
 
         Methods
@@ -328,7 +328,7 @@ def __init__(
             polynomial_degree: int = 4,
             method: str = "legendre",
             defects_type: DefectType = DefectType.EXPLICIT,
-            add_initial_collocation_point: bool = False,
+            include_starting_collocation_point: bool = False,
         ):
             """
             Parameters
@@ -339,8 +339,8 @@ def __init__(
 
             super(OdeSolver.COLLOCATION, self).__init__()
             self.polynomial_degree = polynomial_degree
-            self.add_initial_collocation_point = add_initial_collocation_point
-            self.n_cx = polynomial_degree + 3 if add_initial_collocation_point else polynomial_degree + 2
+            self.include_starting_collocation_point = include_starting_collocation_point
+            self.n_cx = polynomial_degree + 3 if include_starting_collocation_point else polynomial_degree + 2
             self.rk_integrator = COLLOCATION
             self.method = method
             self.defects_type = defects_type
@@ -366,7 +366,7 @@ def integrator(self, ocp, nlp, dynamics_index: int, node_index: int) -> list:
                     "developers and ping @EveCharbie"
                 )
 
-            if self.add_initial_collocation_point:
+            if self.include_starting_collocation_point:
                 x_unscaled = (nlp.states.cx_intermediates_list,)
                 x_scaled = nlp.states.scaled.cx_intermediates_list
             else:

bioptim/examples/stochastic_optimal_control/arm_reaching_muscle_driven.py

@@ -88,7 +88,7 @@ def stochastic_forward_dynamics(
     if with_noise:
         ref = DynamicsFunctions.get(nlp.stochastic_variables["ref"], stochastic_variables)
         k = DynamicsFunctions.get(nlp.stochastic_variables["k"], stochastic_variables)
-        k_matrix = StochasticBioModel.reshape_sym_to_matrix(k, nlp.model.matrix_shape_k)
+        k_matrix = StochasticBioModel.reshape_to_matrix(k, nlp.model.matrix_shape_k)
 
         hand_pos_velo = nlp.model.sensory_reference(states, controls, parameters, stochastic_variables, nlp)
 
@@ -186,7 +186,7 @@ def get_cov_mat(nlp, node_index):
 
     sigma_w = cas.vertcat(nlp.model.sensory_noise_sym, nlp.model.motor_noise_sym) * cas.MX_eye(6)
     cov_sym = cas.MX.sym("cov", nlp.integrated_values.cx_start.shape[0])
-    cov_matrix = StochasticBioModel.reshape_sym_to_matrix(cov_sym, nlp.model.matrix_shape_cov)
+    cov_matrix = StochasticBioModel.reshape_to_matrix(cov_sym, nlp.model.matrix_shape_cov)
 
     dx = stochastic_forward_dynamics(
         nlp.states.cx_start,
@@ -239,7 +239,7 @@ def reach_target_consistantly(controllers: list[PenaltyController]) -> cas.MX:
     q_sym = cas.MX.sym("q_sym", controllers[-1].states["q"].cx_start.shape[0])
     qdot_sym = cas.MX.sym("qdot_sym", controllers[-1].states["qdot"].cx_start.shape[0])
     cov_sym = cas.MX.sym("cov", controllers[-1].integrated_values.cx_start.shape[0])
-    cov_matrix = StochasticBioModel.reshape_sym_to_matrix(cov_sym, controllers[-1].model.matrix_shape_cov)
+    cov_matrix = StochasticBioModel.reshape_to_matrix(cov_sym, controllers[-1].model.matrix_shape_cov)
 
     hand_pos = controllers[0].model.end_effector_position(q_sym)
     hand_vel = controllers[0].model.end_effector_velocity(q_sym, qdot_sym)
@@ -287,10 +287,10 @@ def expected_feedback_effort(controllers: list[PenaltyController], sensory_noise
     # Get the symbolic variables
     ref = controllers[0].stochastic_variables["ref"].cx_start
     cov_sym = cas.MX.sym("cov", controllers[0].integrated_values.cx_start.shape[0])
-    cov_matrix = StochasticBioModel.reshape_sym_to_matrix(cov_sym, controllers[0].model.matrix_shape_cov)
+    cov_matrix = StochasticBioModel.reshape_to_matrix(cov_sym, controllers[0].model.matrix_shape_cov)
 
     k = controllers[0].stochastic_variables["k"].cx_start
-    k_matrix = StochasticBioModel.reshape_sym_to_matrix(k, controllers[0].model.matrix_shape_k)
+    k_matrix = StochasticBioModel.reshape_to_matrix(k, controllers[0].model.matrix_shape_k)
 
     # Compute the expected effort
     hand_pos_velo = controllers[0].model.sensory_reference(

bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_collocations.py

@@ -167,7 +167,6 @@ def prepare_socp(
     dynamics.add(
         DynamicsFcn.STOCHASTIC_TORQUE_DRIVEN,
         problem_type=problem_type,
-        with_cholesky=False,
         expand_dynamics=True,
     )
 

bioptim/examples/stochastic_optimal_control/arm_reaching_torque_driven_explicit.py

@@ -180,7 +180,7 @@ def get_cov_mat(nlp, node_index):
         cas.vertcat(nlp.model.sensory_noise_sym, nlp.model.motor_noise_sym).shape[0]
     )
     cov_sym = cas.MX.sym("cov", nlp.integrated_values.cx_start.shape[0])
-    cov_matrix = StochasticBioModel.reshape_sym_to_matrix(cov_sym, nlp.model.matrix_shape_cov)
+    cov_matrix = StochasticBioModel.reshape_to_matrix(cov_sym, nlp.model.matrix_shape_cov)
 
     dx = stochastic_forward_dynamics(
         nlp.states.cx_start,
@@ -239,7 +239,7 @@ def reach_target_consistently(controllers: list[PenaltyController]) -> cas.MX:
     q_sym = cas.MX.sym("q_sym", controllers[-1].states["q"].cx_start.shape[0])
     qdot_sym = cas.MX.sym("qdot_sym", controllers[-1].states["qdot"].cx_start.shape[0])
     cov_sym = cas.MX.sym("cov", controllers[-1].integrated_values.cx_start.shape[0])
-    cov_matrix = StochasticBioModel.reshape_sym_to_matrix(cov_sym, controllers[-1].model.matrix_shape_cov)
+    cov_matrix = StochasticBioModel.reshape_to_matrix(cov_sym, controllers[-1].model.matrix_shape_cov)
 
     hand_pos = controllers[0].model.markers(q_sym)[2][:2]
     hand_vel = controllers[0].model.marker_velocities(q_sym, qdot_sym)[2][:2]
@@ -287,10 +287,10 @@ def expected_feedback_effort(controllers: list[PenaltyController]) -> cas.MX:
     # Get the symbolic variables
     ref = controllers[0].stochastic_variables["ref"].cx_start
     cov_sym = cas.MX.sym("cov", controllers[0].integrated_values.cx_start.shape[0])
-    cov_matrix = StochasticBioModel.reshape_sym_to_matrix(cov_sym, controllers[0].model.matrix_shape_cov)
+    cov_matrix = StochasticBioModel.reshape_to_matrix(cov_sym, controllers[0].model.matrix_shape_cov)
 
     k = controllers[0].stochastic_variables["k"].cx_start
-    k_matrix = StochasticBioModel.reshape_sym_to_matrix(k, controllers[0].model.matrix_shape_k)
+    k_matrix = StochasticBioModel.reshape_to_matrix(k, controllers[0].model.matrix_shape_k)
 
     # Compute the expected effort
     trace_k_sensor_k = cas.trace(k_matrix @ sensory_noise_matrix @ k_matrix.T)

bioptim/examples/stochastic_optimal_control/common.py

@@ -21,7 +21,7 @@ def dynamics_torque_driven_with_feedbacks(states, controls, parameters, stochast
     if with_noise:
         ref = DynamicsFunctions.get(nlp.stochastic_variables["ref"], stochastic_variables)
         k = DynamicsFunctions.get(nlp.stochastic_variables["k"], stochastic_variables)
-        k_matrix = StochasticBioModel.reshape_sym_to_matrix(k, nlp.model.matrix_shape_k)
+        k_matrix = StochasticBioModel.reshape_to_matrix(k, nlp.model.matrix_shape_k)
 
         motor_noise = nlp.model.motor_noise_sym
         sensory_noise = nlp.model.sensory_noise_sym

bioptim/examples/stochastic_optimal_control/obstacle_avoidance_direct_collocation.py

@@ -298,7 +298,7 @@ def prepare_socp(
             expand_dynamics=expand_dynamics,
         )
         ode_solver = OdeSolver.COLLOCATION(
-            polynomial_degree=socp_type.polynomial_degree, method=socp_type.method, add_initial_collocation_point=True
+            polynomial_degree=socp_type.polynomial_degree, method=socp_type.method, include_starting_collocation_point=True
         )
 
         return OptimalControlProgram(

bioptim/gui/plot.py

@@ -919,7 +919,7 @@ def update_data(self, v: dict):
                                     )
                                 else:
                                     if (
-                                        self.plot_func[key][i].label == "CONTINUITY"
+                                        self.plot_func[key][i].label == "STATE_CONTINUITY"
                                         and nlp.ode_solver.is_direct_collocation
                                     ):
                                         states = state[:, node_idx * (step_size) : (node_idx + 1) * (step_size) + 1]

bioptim/interfaces/stochastic_bio_model.py

@@ -46,20 +46,6 @@ def reshape_to_matrix(var, shape):
                 matrix[s1, s0] = var[s0 * shape_0 + s1]
         return matrix
 
-    @staticmethod
-    def reshape_sym_to_matrix(var, shape):
-        """
-        Restore the matrix form of the variables
-        """
-
-        shape_0, shape_1 = shape
-        matrix = type(var).zeros(shape_0, shape_1)
-
-        for s0 in range(shape_1):
-            for s1 in range(shape_0):
-                matrix[s1, s0] = var[s0 * shape_0 + s1]
-        return matrix
-
     @staticmethod
     def reshape_to_cholesky_matrix(var, shape):
         """
@@ -75,20 +61,6 @@ def reshape_to_cholesky_matrix(var, shape):
                 i += 1
         return matrix
 
-    @staticmethod
-    def reshape_sym_to_cholesky_matrix(var, shape):
-        """
-        Restore the lower diagonal matrix form of the variables vector
-        """
-        shape_0, _ = shape
-        matrix = type(var).zeros(shape_0, shape_0)
-        i = 0
-        for s0 in range(shape_0):
-            for s1 in range(s0 + 1):
-                matrix[s1, s0] = var[i]
-                i += 1
-        return matrix
-
     @staticmethod
     def reshape_to_vector(matrix):
         """

bioptim/limits/constraints.py

@@ -751,7 +751,8 @@ def stochastic_helper_matrix_collocation(
 
             collocation_method = controller.get_nlp.ode_solver.method
             polynomial_degree = controller.get_nlp.ode_solver.polynomial_degree
-            Mc, _ = ConstraintFunction.Functions.collocation_fun_jac(
+            Mc, _ = ConstraintFunction.Functions.collocation_jacobians(
+                penalty,
                 controller,
                 collocation_method,
                 polynomial_degree,
@@ -796,7 +797,8 @@ def stochastic_covariance_matrix_continuity_collocation(
 
             collocation_method = controller.get_nlp.ode_solver.method
             polynomial_degree = controller.get_nlp.ode_solver.polynomial_degree
-            _, Pf = ConstraintFunction.Functions.collocation_fun_jac(
+            _, Pf = ConstraintFunction.Functions.collocation_jacobians(
+                penalty,
                 controller,
                 collocation_method,
                 polynomial_degree,
@@ -896,7 +898,7 @@ def semidefinite_positive_matrix(
             return diagonal_terms
 
         @staticmethod
-        def collocation_fun_jac(controller, method, polynomial_degree):
+        def collocation_jacobians(penalty, controller, method, polynomial_degree):
             def prepare_collocation(method, polynomial_degree):
                 # Get collocation points
                 tau_root = np.append(0, collocation_points(polynomial_degree, method))
@@ -988,16 +990,13 @@ def prepare_collocation(method, polynomial_degree):
             for i in range(polynomial_degree + 1):
                 G_joints = vertcat(G_joints, G_argout[i][joints_index])
 
-            # The function F in x_{k+1} = F(z_k)
-            F = Function("F", [z_q_root, z_q_joints, z_qdot_root, z_qdot_joints], [xf]).expand()
-
             Gdx = jacobian(G_joints, horzcat(x_q_joints, x_qdot_joints))
 
             Gdz = jacobian(G_joints, horzcat(z_q_joints, z_qdot_joints))
 
             Gdw = jacobian(G_joints, vertcat(controller.model.motor_noise_sym, controller.model.sensory_noise_sym))
 
-            Fdz = jacobian(F(z_q_root, z_q_joints, z_qdot_root, z_qdot_joints), horzcat(z_q_joints, z_qdot_joints))
+            Fdz = jacobian(xf, horzcat(z_q_joints, z_qdot_joints))
 
             # Constraint Equality defining M
             Mc = Function(
@@ -1018,7 +1017,9 @@ def prepare_collocation(method, polynomial_degree):
                     controller.model.sensory_noise_sym,
                 ],
                 [Fdz.T - Gdz.T @ m_matrix.T],
-            ).expand()
+            )
+            if penalty.expand:
+                Mc = Mc.expand()
 
             # Covariance propagation rule
             Pf = Function(
@@ -1039,7 +1040,9 @@ def prepare_collocation(method, polynomial_degree):
                     controller.model.sensory_noise_sym,
                 ],
                 [m_matrix @ (Gdx @ cov_matrix @ Gdx.T + Gdw @ sigma_ww @ Gdw.T) @ m_matrix.T],
-            ).expand()
+            )
+            if penalty.expand:
+                Pf = Pf.expand()
 
             return Mc, Pf
 
@@ -1062,9 +1065,9 @@ def get_type() -> Callable
         Returns the type of the penalty
     """
 
-    CONTINUITY = (PenaltyFunctionAbstract.Functions.state_continuity,)
+    STATE_CONTINUITY = (PenaltyFunctionAbstract.Functions.state_continuity,)
     FIRST_COLLOCATION_HELPER_EQUALS_STATE = (
-        PenaltyFunctionAbstract.Functions.first_collocation_helper_equals_final_state,
+        PenaltyFunctionAbstract.Functions.first_collocation_point_equals_state,
     )
     CUSTOM = (PenaltyFunctionAbstract.Functions.custom,)
     NON_SLIPPING = (ConstraintFunction.Functions.non_slipping,)

bioptim/limits/objective_functions.py

@@ -370,7 +370,7 @@ def get_type() -> Callable
             Returns the type of the penalty
         """
 
-        CONTINUITY = (PenaltyFunctionAbstract.Functions.state_continuity,)
+        STATE_CONTINUITY = (PenaltyFunctionAbstract.Functions.state_continuity,)
         CUSTOM = (PenaltyFunctionAbstract.Functions.custom,)
         MINIMIZE_ANGULAR_MOMENTUM = (PenaltyFunctionAbstract.Functions.minimize_angular_momentum,)
         MINIMIZE_COM_ACCELERATION = (PenaltyFunctionAbstract.Functions.minimize_com_acceleration,)

bioptim/limits/penalty.py

@@ -1140,7 +1140,7 @@ def state_continuity(penalty: PenaltyOption, controller: PenaltyController | lis
 
             continuity = controller.states.cx_end
             if controller.get_nlp.ode_solver.is_direct_collocation:
-                if controller.get_nlp.ode_solver.add_initial_collocation_point:
+                if controller.get_nlp.ode_solver.include_starting_collocation_point:
                     cx = horzcat(*controller.states.cx_intermediates_list)
                 else:
                     cx = horzcat(*([controller.states.cx_start] + controller.states.cx_intermediates_list))
@@ -1170,10 +1170,10 @@ def state_continuity(penalty: PenaltyOption, controller: PenaltyController | lis
             return continuity
 
         @staticmethod
-        def first_collocation_helper_equals_final_state(penalty: PenaltyOption, controller: PenaltyController | list):
+        def first_collocation_point_equals_state(penalty: PenaltyOption, controller: PenaltyController | list):
             """
             Insures that the first collocation helper is equal to the states at the shooting node.
-            This is a necessary constraint for COLLOCATION with add_initial_collocation_point.
+            This is a necessary constraint for COLLOCATION with include_starting_collocation_point.
             """
             collocation_helper = controller.states.cx_intermediates_list[0]
             states = controller.states.cx_start

bioptim/optimization/optimal_control_program.py

@@ -949,10 +949,10 @@ def _declare_continuity(self) -> None:
                 # Continuity as constraints
                 if nlp.phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE:
                     penalty = Constraint(
-                        ConstraintFcn.CONTINUITY, node=Node.ALL_SHOOTING, penalty_type=PenaltyType.INTERNAL
+                        ConstraintFcn.STATE_CONTINUITY, node=Node.ALL_SHOOTING, penalty_type=PenaltyType.INTERNAL
                     )
                     penalty.add_or_replace_to_penalty_pool(self, nlp)
-                    if nlp.ode_solver.is_direct_collocation and nlp.ode_solver.add_initial_collocation_point:
+                    if nlp.ode_solver.is_direct_collocation and nlp.ode_solver.include_starting_collocation_point:
                         penalty = Constraint(
                             ConstraintFcn.FIRST_COLLOCATION_HELPER_EQUALS_STATE,
                             node=Node.ALL_SHOOTING,
@@ -962,10 +962,10 @@ def _declare_continuity(self) -> None:
                 else:
                     for shooting_node in range(nlp.ns):
                         penalty = Constraint(
-                            ConstraintFcn.CONTINUITY, node=shooting_node, penalty_type=PenaltyType.INTERNAL
+                            ConstraintFcn.STATE_CONTINUITY, node=shooting_node, penalty_type=PenaltyType.INTERNAL
                         )
                         penalty.add_or_replace_to_penalty_pool(self, nlp)
-                        if nlp.ode_solver.is_direct_collocation and nlp.ode_solver.add_initial_collocation_point:
+                        if nlp.ode_solver.is_direct_collocation and nlp.ode_solver.include_starting_collocation_point:
                             penalty = Constraint(
                                 ConstraintFcn.FIRST_COLLOCATION_HELPER_EQUALS_STATE,
                                 node=shooting_node,
@@ -976,7 +976,7 @@ def _declare_continuity(self) -> None:
                 # Continuity as objectives
                 if nlp.phase_dynamics == PhaseDynamics.SHARED_DURING_THE_PHASE:
                     penalty = Objective(
-                        ObjectiveFcn.Mayer.CONTINUITY,
+                        ObjectiveFcn.Mayer.STATE_CONTINUITY,
                         weight=nlp.dynamics_type.state_continuity_weight,
                         quadratic=True,
                         node=Node.ALL_SHOOTING,
@@ -986,7 +986,7 @@ def _declare_continuity(self) -> None:
                 else:
                     for shooting_point in range(nlp.ns):
                         penalty = Objective(
-                            ObjectiveFcn.Mayer.CONTINUITY,
+                            ObjectiveFcn.Mayer.STATE_CONTINUITY,
                             weight=nlp.dynamics_type.state_continuity_weight,
                             quadratic=True,
                             node=shooting_point,
@@ -1392,7 +1392,7 @@ def compute_penalty_values(t, x, u, p, s, penalty, dt: int | Callable):
                 if not np.all(
                     x == 0
                 ):  # This is a hack to initialize the plots because it x is (N,2) and we need (N, M) in collocation
-                    state_value = x[:, :] if penalty.name == "CONTINUITY" else x[:, [0, -1]]
+                    state_value = x[:, :] if penalty.name == "STATE_CONTINUITY" else x[:, [0, -1]]
                     state_value = state_value.reshape((-1, 1))
                     control_value = control_value.reshape((-1, 1))
                     stochastic_value = stochastic_value.reshape((-1, 1))