Merge pull request #896 from Ipuch/viewer

feat: tracked markers displayed with pyorerun

bioptim/models/biorbd/viewer_pyorerun.py

@@ -1,8 +1,10 @@
 import biorbd_casadi as biorbd
 import numpy as np
 import pyorerun
+from pyomeca import Markers as PyoMarkers
 from typing import Any
 
+from .viewer_utils import _prepare_tracked_markers_for_animation
 from .biorbd_model import BiorbdModel
 from .multi_biorbd_model import MultiBiorbdModel
 from ...optimization.solution.solution_data import SolutionMerge
@@ -41,12 +43,7 @@ def prepare_pyorerun_animation(ocp, solution, show_now=True, show_tracked_marker
             models += [nlp.model.model]
 
     if show_tracked_markers:
-        raise NotImplementedError(
-            "Tracking markers is not implemented for pyorerun. "
-            "Set show_tracked_markers to False such that sol.animate(show_tracked_markers=False)."
-        )
-        # TODO : Implement tracking markers for pyorerun
-        #  with the _prepare_tracked_markers_for_animation in viewer_utils.py
+        tracked_markers = _prepare_tracked_markers_for_animation(ocp.nlp, n_shooting=None)
     else:
         tracked_markers = None
 
@@ -116,11 +113,15 @@ def launch_rerun(
             )
 
         biorbd_model = pyorerun.BiorbdModel.from_biorbd_object(model)
-
+        tm = (
+            PyoMarkers(tm, channels=[n.to_string() for n in biorbd_model.model.markerNames()])
+            if tm is not None
+            else None
+        )
         prerun.add_animated_model(
             biorbd_model,
             data["q"],
-            tracked_markers=tm if tm is not None else None,
+            tracked_markers=tm,
             phase=idx_phase,
         )
 

bioptim/models/biorbd/viewer_utils.py

@@ -13,25 +13,25 @@ def _prepare_tracked_markers_for_animation(
     all_tracked_markers = []
 
     for phase, nlp in enumerate(nlps):
-        n_frames = sum(nlp.ns) + 1 if n_shooting is None else n_shooting + 1
+        n_frames = nlp.ns + 1 if n_shooting is None else n_shooting + 1
 
         n_states_nodes = nlp.n_states_nodes
 
         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):
-                    tracked_markers = np.full((3, nlp.model.nb_markers, n_states_nodes), np.nan)
-
-                    for i in range(len(objective.rows)):
-                        tracked_markers[objective.rows[i], objective.cols, :] = objective.target[i, :, :]
-
-                    missing_row = np.where(np.isnan(tracked_markers))[0]
-                    if missing_row.size > 0:
-                        tracked_markers[missing_row, :, :] = 0
+
+            objective_has_a_target = objective.target is not None
+            objective_is_tracking_markers = objective.type in (
+                ObjectiveFcn.Mayer.TRACK_MARKERS,
+                ObjectiveFcn.Lagrange.TRACK_MARKERS,
+            )
+            objective_is_tracking_all_nodes = objective.node[0] in (Node.ALL, Node.ALL_SHOOTING)
+
+            if objective_has_a_target and objective_is_tracking_markers and objective_is_tracking_all_nodes:
+                tracked_markers = np.zeros((3, nlp.model.nb_markers, n_states_nodes))
+
+                for i, row in enumerate(objective.rows):
+                    tracked_markers[row, objective.cols, :] = objective.target[i, ...]
 
         # interpolation
         if n_frames > 0 and tracked_markers is not None:

tests/shard3/test_global_torque_driven_ocp.py

@@ -428,6 +428,16 @@ def test_track_marker_2D_pendulum(ode_solver, defects_type, phase_dynamics):
         tracked_markers[0][1:, :, -1], np.array([[0.76078505, 0.11005192], [0.98565045, 0.65998405]])
     )
 
+    # testing that preparing tracked markers for animation properly works
+    tracked_markers = _prepare_tracked_markers_for_animation(sol.ocp.nlp, None)
+    npt.assert_equal(tracked_markers[0].shape, (3, 2, n_shooting + 1))
+    npt.assert_equal(tracked_markers[0][0, :, :], np.zeros((2, n_shooting + 1)))
+    npt.assert_almost_equal(tracked_markers[0][1:, :, 0], np.array([[0.82873751, 0.5612772], [0.22793516, 0.24205527]]))
+    npt.assert_almost_equal(tracked_markers[0][1:, :, 5], np.array([[0.80219698, 0.02541913], [0.5107473, 0.36778313]]))
+    npt.assert_almost_equal(
+        tracked_markers[0][1:, :, -1], np.array([[0.76078505, 0.11005192], [0.98565045, 0.65998405]])
+    )
+
 
 @pytest.mark.parametrize("phase_dynamics", [PhaseDynamics.SHARED_DURING_THE_PHASE])
 def test_trampo_quaternions(phase_dynamics):