Merge remote-tracking branch 'origin/cx_intermediates' into cx_intermediates

Author: EveCharbie

bioptim/examples/getting_started/example_external_forces.py

@@ -149,45 +149,45 @@ def setup_external_forces(
     # Add appropriate forces based on method
     if external_force_method == "translational_force":
         external_force_set.add_translational_force(
-            "contact0",
+            "g",
             "Seg1",
             Seg1_force,
             point_of_application_in_local=Seg1_point_of_application if use_point_of_applications else None,
         )
         external_force_set.add_translational_force(
-            "contact0",
+            "h",
             "Test",
             Test_force,
             point_of_application_in_local=Test_point_of_application if use_point_of_applications else None,
         )
 
     elif external_force_method == "translational_force_on_a_marker":
-        external_force_set.add_translational_force("contact0","Test", Test_force, point_of_application_in_local="m0")
+        external_force_set.add_translational_force("q","Test", Test_force, point_of_application_in_local="m0")
 
     elif external_force_method == "in_global":
         external_force_set.add(
-            "contact0",
+            "i",
             "Seg1",
             np.concatenate((Seg1_force, Seg1_force), axis=0),
             point_of_application=Seg1_point_of_application if use_point_of_applications else None,
         )
         external_force_set.add(
-            "contact0",
+            "j",
             "Test",
             np.concatenate((Test_force, Test_force), axis=0),
             point_of_application=Test_point_of_application if use_point_of_applications else None,
         )
     elif external_force_method == "in_global_torque":
-        external_force_set.add_torque("contact0", "Seg1", Seg1_force)
-        external_force_set.add_torque("contact0", "Test", Test_force)
+        external_force_set.add_torque("k", "Seg1", Seg1_force)
+        external_force_set.add_torque("l", "Test", Test_force)
 
     elif external_force_method == "in_segment_torque":
-        external_force_set.add_torque_in_segment_frame("contact0", "Seg1", Seg1_force)
-        external_force_set.add_torque_in_segment_frame("contact0", "Test", Test_force)
+        external_force_set.add_torque_in_segment_frame("m", "Seg1", Seg1_force)
+        external_force_set.add_torque_in_segment_frame("n", "Test", Test_force)
 
     elif external_force_method == "in_segment":
-        external_force_set.add_in_segment_frame("contact0", "Seg1", np.concatenate((Seg1_force, Seg1_force), axis=0))
-        external_force_set.add_in_segment_frame("contact0", "Test", np.concatenate((Test_force, Test_force), axis=0))
+        external_force_set.add_in_segment_frame("o", "Seg1", np.concatenate((Seg1_force, Seg1_force), axis=0))
+        external_force_set.add_in_segment_frame("p", "Test", np.concatenate((Test_force, Test_force), axis=0))
 
     return external_force_set
 

bioptim/examples/muscle_driven_with_contact/contact_forces_inverse_dynamics_constraint_muscle.py

@@ -48,7 +48,12 @@ def custom_configure(ocp: OptimalControlProgram, nlp: NonLinearProgram, numerica
 
     # Implicit variables
     ConfigureProblem.configure_rigid_contact_forces(
-        ocp, nlp, as_states=False, as_algebraic_states=True, as_controls=False
+        ocp,
+        nlp,
+        as_states=False,
+        as_algebraic_states=True,
+        as_controls=False,
+        as_states_dot=False,
     )
 
     # Dynamics
@@ -72,7 +77,6 @@ def custom_dynamics(
     qdot = nlp.get_var_from_states_or_controls("qdot", states, controls)
     residual_tau = nlp.get_var_from_states_or_controls("tau", states, controls)
     mus_activations = nlp.get_var_from_states_or_controls("muscles", states, controls)
-    rigid_contact_forces_derivatives = nlp.get_var_from_states_or_controls()
 
     # Get external forces from the states
     rigid_contact_forces = nlp.get_external_forces(
@@ -158,7 +162,7 @@ def prepare_ocp(biorbd_model_path, phase_time, n_shooting, expand_dynamics=True)
         node=Node.ALL_SHOOTING,
     )
     multinode_constraints = MultinodeConstraintList()
-    for i_node in range(n_shooting):
+    for i_node in range(n_shooting - 1):
         multinode_constraints.add(
             MultinodeConstraintFcn.ALGEBRAIC_STATES_CONTINUITY,
             nodes_phase=(0, 0),

bioptim/interfaces/interface_utils.py

@@ -368,16 +368,22 @@ def _get_weighted_function_inputs(penalty, penalty_idx, ocp, nlp, scaled):
 
     if nlp:
         x = PenaltyHelpers.states(
-            penalty, penalty_idx, lambda p_idx, n_idx, sn_idx: _get_x(ocp, p_idx, n_idx, sn_idx, scaled)
+            penalty,
+            penalty_idx,
+            lambda p_idx, n_idx, sn_idx: _get_x(ocp, p_idx, n_idx, sn_idx, scaled, penalty),
         )
         u = PenaltyHelpers.controls(
-            penalty, penalty_idx, lambda p_idx, n_idx, sn_idx: _get_u(ocp, p_idx, n_idx, sn_idx, scaled)
+            penalty,
+            penalty_idx,
+            lambda p_idx, n_idx, sn_idx: _get_u(ocp, p_idx, n_idx, sn_idx, scaled, penalty),
         )
         p = PenaltyHelpers.parameters(
             penalty, penalty_idx, lambda p_idx, n_idx, sn_idx: _get_p(ocp, p_idx, n_idx, sn_idx, scaled)
         )
         a = PenaltyHelpers.states(
-            penalty, penalty_idx, lambda p_idx, n_idx, sn_idx: _get_a(ocp, p_idx, n_idx, sn_idx, scaled)
+            penalty,
+            penalty_idx,
+            lambda p_idx, n_idx, sn_idx: _get_a(ocp, p_idx, n_idx, sn_idx, scaled, penalty),
         )
         d = PenaltyHelpers.numerical_timeseries(
             penalty,
@@ -396,7 +402,9 @@ def _get_weighted_function_inputs(penalty, penalty_idx, ocp, nlp, scaled):
     return t0, x, u, p, a, d, weight, target
 
 
-def _get_x(ocp, phase_idx, node_idx, subnodes_idx, scaled):
+def _get_x(ocp, phase_idx, node_idx, subnodes_idx, scaled, penalty):
+    idx = 0 if not penalty.is_multinode_penalty else penalty.nodes_phase.index(phase_idx)
+    subnodes_are_decision_states = penalty.subnodes_are_decision_states[idx] and not penalty.is_transition
     values = ocp.nlp[phase_idx].X_scaled if scaled else ocp.nlp[phase_idx].X
     if subnodes_idx.stop == -1:
         if subnodes_idx.start == 0:
@@ -407,11 +415,16 @@ def _get_x(ocp, phase_idx, node_idx, subnodes_idx, scaled):
         else:
             raise RuntimeError("only subnodes_idx.start == 0 is supported for subnodes_idx.stop == -1")
     else:
-        x = values[node_idx][:, subnodes_idx] if node_idx < len(values) else ocp.cx()
+        if subnodes_are_decision_states:
+            x = values[node_idx][:, subnodes_idx] if node_idx < len(values) else ocp.cx()
+        else:
+            x = values[node_idx][:, 0] if node_idx < len(values) else ocp.cx()
     return x
 
 
-def _get_u(ocp, phase_idx, node_idx, subnodes_idx, scaled):
+def _get_u(ocp, phase_idx, node_idx, subnodes_idx, scaled, penalty):
+    idx = 0 if not penalty.is_multinode_penalty else penalty.nodes_phase.index(phase_idx)
+    subnodes_are_decision_states = penalty.subnodes_are_decision_states[idx] and not penalty.is_transition
     values = ocp.nlp[phase_idx].U_scaled if scaled else ocp.nlp[phase_idx].U
     if subnodes_idx.stop == -1:
         if subnodes_idx.start == 0:
@@ -422,15 +435,20 @@ def _get_u(ocp, phase_idx, node_idx, subnodes_idx, scaled):
         else:
             raise RuntimeError("only subnodes_idx.start == 0 is supported for subnodes_idx.stop == -1")
     else:
-        u = values[node_idx][:, subnodes_idx] if node_idx < len(values) else ocp.cx()
+        if subnodes_are_decision_states:
+            u = values[node_idx][:, subnodes_idx] if node_idx < len(values) else ocp.cx()
+        else:
+            u = values[node_idx][:, 0] if node_idx < len(values) else ocp.cx()
     return u
 
 
 def _get_p(ocp, phase_idx, node_idx, subnodes_idx, scaled):
     return ocp.parameters.scaled.cx if scaled else ocp.parameters.scaled
 
 
-def _get_a(ocp, phase_idx, node_idx, subnodes_idx, scaled):
+def _get_a(ocp, phase_idx, node_idx, subnodes_idx, scaled, penalty):
+    idx = 0 if not penalty.is_multinode_penalty else penalty.nodes_phase.index(phase_idx)
+    subnodes_are_decision_states = penalty.subnodes_are_decision_states[idx] and not penalty.is_transition
     values = ocp.nlp[phase_idx].A_scaled if scaled else ocp.nlp[phase_idx].A
     if subnodes_idx.stop == -1:
         if subnodes_idx.start == 0:
@@ -441,7 +459,10 @@ def _get_a(ocp, phase_idx, node_idx, subnodes_idx, scaled):
         else:
             raise RuntimeError("only subnodes_idx.start == 0 is supported for subnodes_idx.stop == -1")
     else:
-        a = values[node_idx][:, subnodes_idx] if node_idx < len(values) else ocp.cx()
+        if subnodes_are_decision_states:
+            a = values[node_idx][:, subnodes_idx] if node_idx < len(values) else ocp.cx()
+        else:
+            a = values[node_idx][:, 0] if node_idx < len(values) else ocp.cx()
     return a
 
 

bioptim/limits/multinode_penalty.py

@@ -68,7 +68,8 @@ def __init__(
         if len(nodes) != len(nodes_phase):
             raise ValueError("Each of the nodes must have a corresponding nodes_phase")
 
-        self.multinode_penalty = True
+        self.is_multinode_penalty = True
+        self.is_transition = False
 
         self.nodes_phase = nodes_phase
         self.nodes = nodes

bioptim/limits/penalty_helpers.py

@@ -29,7 +29,7 @@ def t0(penalty, index, get_t0: Callable):
         This method returns the t0 of a penalty.
         """
 
-        if penalty.multinode_penalty:
+        if penalty.is_multinode_penalty:
             phases, nodes, _ = _get_multinode_indices(penalty, is_constructing_penalty=False)
             phase, node = phases[0], nodes[0]
         else:
@@ -68,20 +68,34 @@ def states(penalty, index, get_state_decision: Callable, is_constructing_penalty
 
         node = penalty.node_idx[index]
 
-        if penalty.multinode_penalty:
+        if penalty.is_multinode_penalty:
             x = []
             phases, nodes, subnodes = _get_multinode_indices(penalty, is_constructing_penalty)
             idx = 0
             for phase, node, sub in zip(phases, nodes, subnodes):
-                if not is_constructing_penalty and node == penalty.ns[idx] and (penalty.control_types[idx] != ControlType.LINEAR_CONTINUOUS and penalty.control_types[idx] != ControlType.CONSTANT_WITH_LAST_NODE):
+                if (
+                    not is_constructing_penalty
+                    and node == penalty.ns[idx]
+                    and (
+                        penalty.control_types[idx] != ControlType.LINEAR_CONTINUOUS
+                        and penalty.control_types[idx] != ControlType.CONSTANT_WITH_LAST_NODE
+                    )
+                ):
                     x.append(_reshape_to_vector(get_state_decision(phase, node, range(0, 1))))
                 else:
                     x.append(_reshape_to_vector(get_state_decision(phase, node, sub)))
                 idx += 1
             return _vertcat(x)
 
         else:
-            subnodes = slice(0, None if node < penalty.ns[0] and penalty.subnodes_are_decision_states[0] else 1)
+            subnodes = slice(
+                0,
+                (
+                    None
+                    if node < penalty.ns[0] and penalty.subnodes_are_decision_states[0] and not penalty.is_transition
+                    else 1
+                ),
+            )
             x0 = _reshape_to_vector(get_state_decision(penalty.phase, node, subnodes))
 
             if is_constructing_penalty:
@@ -97,12 +111,23 @@ def states(penalty, index, get_state_decision: Callable, is_constructing_penalty
     def controls(penalty, index, get_control_decision: Callable, is_constructing_penalty: bool = False):
         node = penalty.node_idx[index]
 
-        if penalty.multinode_penalty:
+        if penalty.is_multinode_penalty:
             u = []
             phases, nodes, subnodes = _get_multinode_indices(penalty, is_constructing_penalty)
+            idx = 0
             for phase, node, sub in zip(phases, nodes, subnodes):
-                # No need to test for control types as this is never integrated (so we only need the starting value)
-                u.append(_reshape_to_vector(get_control_decision(phase, node, sub)))
+                if (
+                    not is_constructing_penalty
+                    and node == penalty.ns[idx]
+                    and (
+                        penalty.control_types[idx] != ControlType.LINEAR_CONTINUOUS
+                        and penalty.control_types[idx] != ControlType.CONSTANT_WITH_LAST_NODE
+                    )
+                ):
+                    u.append(_reshape_to_vector(get_control_decision(phase, node, range(0, 1))))
+                else:
+                    u.append(_reshape_to_vector(get_control_decision(phase, node, sub)))
+                idx += 1
             return _vertcat(u)
 
         if is_constructing_penalty:
@@ -136,7 +161,7 @@ def parameters(penalty, index, get_parameter_decision: Callable):
     @staticmethod
     def numerical_timeseries(penalty, index, get_numerical_timeseries: Callable):
         node = penalty.node_idx[index]
-        if penalty.multinode_penalty:
+        if penalty.is_multinode_penalty:
             # numerical timeseries are expected to be provided only at the shooting node.
             for i_phase in penalty.nodes_phase:
                 d = get_numerical_timeseries(i_phase, node, 0)  # cx_start
@@ -219,27 +244,26 @@ def get_multinode_penalty_subnodes_starting_index(p):
 
 
 def _get_multinode_indices(penalty, is_constructing_penalty: bool):
-    if not penalty.multinode_penalty:
+    if not penalty.is_multinode_penalty:
         raise RuntimeError("This function should only be called for multinode penalties")
 
-    if not (all(penalty.subnodes_are_decision_states) or sum(penalty.subnodes_are_decision_states) == 0):
-        # This check allows to test only for penalty.subnodes_are_decision_states[0] below
-        raise NotImplementedError(
-            "All controllers must be of the same type (either all or none should have subnodes_are_decision_states)"
-        )
-
     phases = penalty.nodes_phase
     nodes = penalty.multinode_idx
 
     startings = PenaltyHelpers.get_multinode_penalty_subnodes_starting_index(penalty)
     subnodes = []
     for i_starting, starting in enumerate(startings):
-        if starting < 0 or starting == 2:  # The last cx accessible
+        if starting < 0:  # The last cx accessible (cx_end)
             if is_constructing_penalty:
                 subnodes.append(slice(-1, None))
             else:
                 subnodes.append(slice(0, 1))
-        elif penalty.subnodes_are_decision_states[0]:
+        elif starting == 2:  # Also the last cx accessible (cx_end) since there are only 3 cx available
+            if is_constructing_penalty:
+                subnodes.append(slice(2, 3))
+            else:
+                subnodes.append(slice(0, 1))
+        elif penalty.subnodes_are_decision_states[i_starting] and not penalty.is_transition:
             if nodes[i_starting] >= penalty.ns[i_starting]:
                 subnodes.append(slice(0, 1))
             else:

bioptim/limits/penalty_option.py

@@ -196,7 +196,8 @@ def __init__(
         self.weighted_function: list[Function | None] = []
         self.weighted_function_non_threaded: list[Function | None] = []
 
-        self.multinode_penalty = False
+        self.is_multinode_penalty = False
+        self.is_transition = False
         self.nodes_phase = None  # This is relevant for multinodes
         self.nodes = None  # This is relevant for multinodes
         if self.derivative and self.explicit_derivative:
@@ -606,9 +607,9 @@ def _set_penalty_function(self, controllers: list[PenaltyController], fcn: MX |
             self.weighted_function[node] = self.weighted_function[node].expand()
 
     def _check_sanity_of_penalty_interactions(self, controller):
-        if self.multinode_penalty and self.explicit_derivative:
+        if self.is_multinode_penalty and self.explicit_derivative:
             raise ValueError("multinode_penalty and explicit_derivative cannot be true simultaneously")
-        if self.multinode_penalty and self.derivative:
+        if self.is_multinode_penalty and self.derivative:
             raise ValueError("multinode_penalty and derivative cannot be true simultaneously")
         if self.derivative and self.explicit_derivative:
             raise ValueError("derivative and explicit_derivative cannot be true simultaneously")
@@ -625,7 +626,7 @@ def _check_sanity_of_penalty_interactions(self, controller):
             )
 
     def get_variable_inputs(self, controllers: list[PenaltyController]):
-        if self.multinode_penalty:
+        if self.is_multinode_penalty:
             controller = controllers[0]  # Recast controller as a normal variable (instead of a list)
             self.node_idx[0] = controller.node_index
 
@@ -733,7 +734,7 @@ def vertcat_cx_end():
                     # performing some kind of integration or derivative and this last node does not exist
                     if nlp.control_type in (ControlType.CONSTANT_WITH_LAST_NODE,):
                         return vertcat(u, controls.scaled.cx_end)
-                    if self.integrate or self.derivative or self.explicit_derivative or self.multinode_penalty:
+                    if self.integrate or self.derivative or self.explicit_derivative or self.is_multinode_penalty:
                         return u
                     else:
                         return vertcat(u, controls.scaled.cx_end)

bioptim/limits/phase_transition.py

@@ -71,6 +71,7 @@ def __init__(
         self.bounds = Bounds("phase_transition", interpolation=InterpolationType.CONSTANT)
         self.node = Node.TRANSITION
         self.quadratic = True
+        self.is_transition = True
 
     def add_or_replace_to_penalty_pool(self, ocp, nlp):
         super(PhaseTransition, self).add_or_replace_to_penalty_pool(ocp, nlp)

bioptim/models/biorbd/external_forces.py

@@ -4,11 +4,6 @@
 
 class ExternalForceSet:
 
-    @property
-    def nb_external_forces(self) -> int:
-        attributes = ["in_global", "torque_in_global", "translational_in_global", "in_local", "torque_in_local"]
-        return sum([len(values) for attr in attributes for values in getattr(self, attr).values()])
-
     @property
     def nb_external_forces_components(self) -> int:
         """Return the number of vertical components of the external forces if concatenated in a unique vector"""
@@ -38,7 +33,7 @@ def check_segment_names(self, segment_names: tuple[str, ...]) -> None:
         for attr in attributes:
             for force_name , force in getattr(self, attr).items():
                 if force["segment"] not in segment_names:
-                    wrong_segments.append(force[0]["segment"])
+                    wrong_segments.append(force["segment"])
 
         if wrong_segments:
             raise ValueError(