Merge pull request #714 from Ipuch/master

some little test on preparing the data for animation of tracked markers.

bioptim/examples/muscle_driven_with_contact/contact_forces_inequality_constraint_muscle.py

@@ -32,7 +32,7 @@ def prepare_ocp(biorbd_model_path, phase_time, n_shooting, min_bound, max_bound)
     tau_min, tau_max, tau_init = -500.0, 500.0, 0.0
     activation_min, activation_max, activation_init = 0.0, 1.0, 0.5
     dof_mapping = BiMappingList()
-    dof_mapping.add("tau", [None, None, None, 0], [3])
+    dof_mapping.add("tau", bimapping=None, to_second=[None, None, None, 0], to_first=[3])
 
     # Add objective functions
     objective_functions = ObjectiveList()
@@ -66,34 +66,36 @@ def prepare_ocp(biorbd_model_path, phase_time, n_shooting, min_bound, max_bound)
 
     # Initialize x_bounds
     x_bounds = BoundsList()
-    x_bounds.add(bounds=bio_model.bounds_from_ranges(["q", "qdot"]))
-    x_bounds[0][:, 0] = pose_at_first_node + [0] * n_qdot
+    x_bounds["q"] = bio_model.bounds_from_ranges("q")
+    x_bounds["q"][:, 0] = pose_at_first_node
+    x_bounds["qdot"] = bio_model.bounds_from_ranges("qdot")
+    x_bounds["qdot"][:, 0] = 0
 
     # Initial guess
     x_init = InitialGuessList()
-    x_init.add(pose_at_first_node + [0] * n_qdot)
+    x_init["q"] = pose_at_first_node
+    x_init["qdot"] = np.zeros((n_qdot,))
 
     # Define control path constraint
     u_bounds = BoundsList()
-    u_bounds.add(
-        [tau_min] * len(dof_mapping["tau"].to_first) + [activation_min] * bio_model.nb_muscles,
-        [tau_max] * len(dof_mapping["tau"].to_first) + [activation_max] * bio_model.nb_muscles,
-    )
+    u_bounds["tau"] = [tau_min] * len(dof_mapping["tau"].to_first), [tau_max] * len(dof_mapping["tau"].to_first)
+    u_bounds["muscles"] = [activation_min] * bio_model.nb_muscles, [activation_max] * bio_model.nb_muscles
 
     u_init = InitialGuessList()
-    u_init.add([tau_init] * len(dof_mapping["tau"].to_first) + [activation_init] * bio_model.nb_muscles)
+    u_init["tau"] = [tau_init] * len(dof_mapping["tau"].to_first)
+    u_init["muscles"] = [activation_init] * bio_model.nb_muscles
     # ------------- #
 
     return OptimalControlProgram(
         bio_model,
         dynamics,
         n_shooting,
         phase_time,
-        x_init,
-        u_init,
-        x_bounds,
-        u_bounds,
-        objective_functions,
+        x_bounds=x_bounds,
+        u_bounds=u_bounds,
+        x_init=x_init,
+        u_init=u_init,
+        objective_functions=objective_functions,
         constraints=constraints,
         variable_mappings=dof_mapping,
         assume_phase_dynamics=True,

bioptim/examples/muscle_driven_with_contact/contact_forces_inequality_constraint_muscle_excitations.py

@@ -22,7 +22,6 @@
     SelectionMapping,
     Dependency,
     BoundsList,
-    Bounds,
     InitialGuessList,
     OdeSolver,
     Solver,
@@ -38,12 +37,9 @@ def prepare_ocp(biorbd_model_path, phase_time, n_shooting, min_bound, ode_solver
     # adds a bimapping to bimappinglist
     # dof_mapping.add("tau", [None, None, None, 0], [3])
     # easier way is to use SelectionMapping which is a subclass of biMapping
-    bimap = SelectionMapping(
-        nb_elements=bio_model.nb_dof,
-        independent_indices=(3,),
-        dependencies=(Dependency(dependent_index=None, reference_index=None, factor=None),),
-    )
-    dof_mapping.add(name="tau", bimapping=bimap)
+    dof_mapping = BiMappingList()
+    dof_mapping.add("tau", bimapping=None, to_second=[None, None, None, 0], to_first=[3])
+
     # Add objective functions
     objective_functions = ObjectiveList()
     objective_functions.add(ObjectiveFcn.Mayer.MINIMIZE_PREDICTED_COM_HEIGHT, weight=-1)
@@ -77,34 +73,38 @@ def prepare_ocp(biorbd_model_path, phase_time, n_shooting, min_bound, ode_solver
 
     # Initialize x_bounds
     x_bounds = BoundsList()
-    x_bounds.add(bounds=bio_model.bounds_from_ranges(["q", "qdot"]))
-    x_bounds[0].concatenate(Bounds([activation_min] * n_mus, [activation_max] * n_mus))
-    x_bounds[0][:, 0] = pose_at_first_node + [0] * n_qdot + [0.5] * n_mus
+    x_bounds["q"] = bio_model.bounds_from_ranges("q")
+    x_bounds["q"][:, 0] = pose_at_first_node
+    x_bounds["qdot"] = bio_model.bounds_from_ranges("qdot")
+    x_bounds["qdot"][:, 0] = np.zeros((n_qdot,))
+    x_bounds["muscles"] = [[activation_min] * n_mus, [activation_max] * n_mus]
+    x_bounds["muscles"][:, 0] = np.zeros((n_mus,))
 
     # Initial guess
     x_init = InitialGuessList()
-    x_init.add(pose_at_first_node + [0] * n_qdot + [0.5] * n_mus)
+    x_init["q"] = pose_at_first_node
+    x_init["qdot"] = np.zeros((n_qdot,))
+    x_init["muscles"] = np.zeros((n_mus,))
 
     # Define control path constraint
     u_bounds = BoundsList()
-    u_bounds.add(
-        [torque_min] * len(dof_mapping["tau"].to_first) + [activation_min] * bio_model.nb_muscles,
-        [torque_max] * len(dof_mapping["tau"].to_first) + [activation_max] * bio_model.nb_muscles,
-    )
+    u_bounds["tau"] = [torque_min] * len(dof_mapping["tau"].to_first), [torque_max] * len(dof_mapping["tau"].to_first)
+    u_bounds["muscles"] = [activation_min] * bio_model.nb_muscles, [activation_max] * bio_model.nb_muscles
 
     u_init = InitialGuessList()
-    u_init.add([torque_init] * len(dof_mapping["tau"].to_first) + [activation_init] * bio_model.nb_muscles)
+    u_init["tau"] = [torque_init] * len(dof_mapping["tau"].to_first)
+    u_init["muscles"] = [activation_init] * bio_model.nb_muscles
     # ------------- #
 
     return OptimalControlProgram(
         bio_model,
         dynamics,
         n_shooting,
         phase_time,
-        x_init,
-        u_init,
-        x_bounds,
-        u_bounds,
+        x_bounds=x_bounds,
+        u_bounds=u_bounds,
+        x_init=x_init,
+        u_init=u_init,
         objective_functions=objective_functions,
         constraints=constraints,
         variable_mappings=dof_mapping,

bioptim/interfaces/biomodel.py

@@ -1,4 +1,5 @@
-from typing import Protocol, Callable
+import numpy as np
+from typing import Protocol, Callable, Any
 
 from casadi import MX, SX
 from ..misc.mapping import BiMapping, BiMappingList
@@ -303,3 +304,26 @@ def partitioned_forward_dynamics(
         ROBOTRAN: a powerful symbolic gnerator of multibody models, Mech. Sci., 4, 199–219,
         https://doi.org/10.5194/ms-4-199-2013, 2013.
         """
+
+    @staticmethod
+    def animate(
+        solution: Any, show_now: bool = True, tracked_markers: list[np.ndarray, ...] = None, **kwargs: Any
+    ) -> None | list:
+        """
+        Animate a solution
+
+        Parameters
+        ----------
+        solution: Any
+            The solution to animate
+        show_now: bool
+            If the animation should be shown immediately or not
+        tracked_markers: list[np.ndarray, ...]
+            The tracked markers (3, n_markers, n_frames)
+        kwargs: dict
+            The options to pass to the animator
+
+        Returns
+        -------
+        The animator object or None if show_now
+        """

bioptim/interfaces/biorbd_model.py

@@ -8,6 +8,7 @@
     GeneralizedAcceleration,
 )
 from casadi import SX, MX, vertcat, horzcat, norm_fro
+import numpy as np
 
 from ..misc.utils import check_version
 from ..limits.path_conditions import Bounds
@@ -512,6 +513,53 @@ def lagrangian(self, q: MX | SX, qdot: MX | SX) -> MX | SX:
         qdot_biorbd = GeneralizedVelocity(qdot)
         return self.model.Lagrangian(q_biorbd, qdot_biorbd).to_mx()
 
+    @staticmethod
+    def animate(
+        solution: Any, show_now: bool = True, tracked_markers: list[np.ndarray, ...] = None, **kwargs: Any
+    ) -> None | list:
+        try:
+            import bioviz
+        except ModuleNotFoundError:
+            raise RuntimeError("bioviz must be install to animate the model")
+
+        check_version(bioviz, "2.3.0", "2.4.0")
+
+        states = solution.states
+        if not isinstance(states, (list, tuple)):
+            states = [states]
+
+        if tracked_markers is None:
+            tracked_markers = [None] * len(states)
+
+        all_bioviz = []
+        for idx_phase, data in enumerate(states):
+            if not isinstance(solution.ocp.nlp[idx_phase].model, BiorbdModel):
+                raise NotImplementedError("Animation is only implemented for biorbd models")
+
+            # This calls each of the function that modify the internal dynamic model based on the parameters
+            nlp = solution.ocp.nlp[idx_phase]
+
+            # noinspection PyTypeChecker
+            biorbd_model: BiorbdModel = nlp.model
+
+            all_bioviz.append(bioviz.Viz(biorbd_model.path, **kwargs))
+            all_bioviz[-1].load_movement(solution.ocp.nlp[idx_phase].variable_mappings["q"].to_second.map(data["q"]))
+
+            if tracked_markers[idx_phase] is not None:
+                all_bioviz[-1].load_experimental_markers(tracked_markers[idx_phase])
+
+        if show_now:
+            b_is_visible = [True] * len(all_bioviz)
+            while sum(b_is_visible):
+                for i, b in enumerate(all_bioviz):
+                    if b.vtk_window.is_active:
+                        b.update()
+                    else:
+                        b_is_visible[i] = False
+            return None
+        else:
+            return all_bioviz
+
 
 class MultiBiorbdModel:
     """
@@ -1090,3 +1138,7 @@ def _qddot_mapping(self, mapping: BiMapping = None) -> dict:
 
     def lagrangian(self):
         raise NotImplementedError("lagrangian is not implemented yet for MultiBiorbdModel")
+
+    @staticmethod
+    def animate(solution: Any, show_now: bool = True, tracked_markers: list = [], **kwargs: Any) -> None | list:
+        raise NotImplementedError("animate is not implemented yet for MultiBiorbdModel")

bioptim/optimization/solution.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 from scipy import interpolate as sci_interp
+from scipy.interpolate import interp1d
 from casadi import vertcat, DM, Function
 from matplotlib import pyplot as plt
 
@@ -24,6 +25,7 @@
 from ..optimization.optimization_vector import OptimizationVectorHelper
 from ..dynamics.ode_solver import OdeSolver
 from ..interfaces.solve_ivp_interface import solve_ivp_interface, solve_ivp_bioptim_interface
+from ..interfaces.biorbd_model import BiorbdModel
 
 
 class Solution:
@@ -1350,7 +1352,12 @@ def graphs(
             plt.show()
 
     def animate(
-        self, n_frames: int = 0, shooting_type: Shooting = None, show_now: bool = True, **kwargs: Any
+        self,
+        n_frames: int = 0,
+        shooting_type: Shooting = None,
+        show_now: bool = True,
+        show_tracked_markers: bool = False,
+        **kwargs: Any,
     ) -> None | list:
         """
         Animate the simulation
@@ -1364,6 +1371,8 @@ def animate(
             The Shooting type to animate
         show_now: bool
             If the bioviz exec() function should be called automatically. This is blocking method
+        show_tracked_markers: bool
+            If the tracked markers should be displayed
         kwargs: Any
             Any parameters to pass to bioviz
 
@@ -1372,61 +1381,82 @@ def animate(
             A list of bioviz structures (one for each phase). So one can call exec() by hand
         """
 
-        try:
-            import bioviz
-        except ModuleNotFoundError:
-            raise RuntimeError("bioviz must be install to animate the model")
-
-        from ..interfaces.biorbd_model import BiorbdModel
-
-        check_version(bioviz, "2.3.0", "2.4.0")
-
         data_to_animate = self.integrate(shooting_type=shooting_type) if shooting_type else self.copy()
         if n_frames == 0:
             try:
-                data_to_animate = data_to_animate.interpolate(sum(self.ns))
+                data_to_animate = data_to_animate.interpolate(sum(self.ns) + 1)
             except RuntimeError:
                 pass
 
         elif n_frames > 0:
             data_to_animate = data_to_animate.interpolate(n_frames)
 
-        states = data_to_animate.states
-        if not isinstance(states, (list, tuple)):
-            states = [states]
+        if show_tracked_markers and len(self.ocp.nlp) == 1:
+            tracked_markers = self._prepare_tracked_markers_for_animation(n_shooting=n_frames)
+        elif show_tracked_markers and len(self.ocp.nlp) > 1:
+            raise NotImplementedError(
+                "Tracking markers is not implemented for multiple phases. "
+                "Set show_tracked_markers to False such that sol.animate(show_tracked_markers=False)."
+            )
+        else:
+            tracked_markers = [None for _ in range(len(self.ocp.nlp))]
+
+        # assuming that all the models or the same type.
+        self._check_models_comes_from_same_super_class()
+        self.ocp.nlp[0].model.animate(
+            solution=data_to_animate,
+            show_now=show_now,
+            tracked_markers=tracked_markers,
+            **kwargs,
+        )
+
+    def _check_models_comes_from_same_super_class(self):
+        """Check that all the models comes from the same super class"""
+        for i, nlp in enumerate(self.ocp.nlp):
+            model_super_classes = nlp.model.__class__.mro()[:-1]  # remove object class
+            nlps = self.ocp.nlp.copy()
+            del nlps[i]
+            for j, sub_nlp in enumerate(nlps):
+                if not any([isinstance(sub_nlp.model, super_class) for super_class in model_super_classes]):
+                    raise RuntimeError(
+                        f"The animation is only available for compatible models. "
+                        f"Here, the model of phase {i} is of type {nlp.model.__class__.__name__} and the model of "
+                        f"phase {j + 1 if i < j else j} is of type {sub_nlp.model.__class__.__name__} and "
+                        f"they don't share the same super class."
+                    )
 
-        all_bioviz = []
-        for idx_phase, data in enumerate(states):
-            if not isinstance(self.ocp.nlp[idx_phase].model, BiorbdModel):
-                raise NotImplementedError("Animation is only implemented for biorbd models")
+    def _prepare_tracked_markers_for_animation(self, n_shooting: int = None) -> list:
+        """Prepare the markers which are tracked to the animation"""
 
-            # This calls each of the function that modify the internal dynamic model based on the parameters
-            nlp = self.ocp.nlp[idx_phase]
+        n_frames = sum(self.ns) + 1 if n_shooting is None else n_shooting + 1
 
-            # noinspection PyTypeChecker
-            biorbd_model: BiorbdModel = nlp.model
+        all_tracked_markers = []
 
-            all_bioviz.append(bioviz.Viz(biorbd_model.path, **kwargs))
-            all_bioviz[-1].load_movement(self.ocp.nlp[idx_phase].variable_mappings["q"].to_second.map(data["q"]))
-            for objective in self.ocp.nlp[idx_phase].J:
+        for phase, nlp in enumerate(self.ocp.nlp):
+            tracked_markers = None
+            for objective in nlp.J:
                 if objective.target is not None:
                     if objective.type in (
                         ObjectiveFcn.Mayer.TRACK_MARKERS,
                         ObjectiveFcn.Lagrange.TRACK_MARKERS,
                     ) and objective.node[0] in (Node.ALL, Node.ALL_SHOOTING):
-                        all_bioviz[-1].load_experimental_markers(objective.target[0])
-
-        if show_now:
-            b_is_visible = [True] * len(all_bioviz)
-            while sum(b_is_visible):
-                for i, b in enumerate(all_bioviz):
-                    if b.vtk_window.is_active:
-                        b.update()
-                    else:
-                        b_is_visible[i] = False
-            return None
-        else:
-            return all_bioviz
+                        tracked_markers = np.full((3, nlp.model.nb_markers, self.ns[phase] + 1), np.nan)
+                        for i in range(len(objective.rows)):
+                            tracked_markers[objective.rows[i], objective.cols, :] = objective.target[0][i, :, :]
+                        missing_row = np.where(np.isnan(tracked_markers))[0]
+                        if missing_row.size > 0:
+                            tracked_markers[missing_row, :, :] = 0
+
+            # interpolation
+            if n_frames > 0 and tracked_markers is not None:
+                x = np.linspace(0, self.ns[phase], self.ns[phase] + 1)
+                xnew = np.linspace(0, self.ns[phase], n_frames)
+                f = interp1d(x, tracked_markers, kind="cubic")
+                tracked_markers = f(xnew)
+
+            all_tracked_markers.append(tracked_markers)
+
+        return all_tracked_markers
 
     def _get_penalty_cost(self, nlp, penalty):
         phase_idx = nlp.phase_idx