diff --git a/bioptim/models/biorbd/viewer_pyorerun.py b/bioptim/models/biorbd/viewer_pyorerun.py index 00d03bd0d..35481ffde 100644 --- a/bioptim/models/biorbd/viewer_pyorerun.py +++ b/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, ) diff --git a/bioptim/models/biorbd/viewer_utils.py b/bioptim/models/biorbd/viewer_utils.py index 067ae73e8..e6a79180e 100644 --- a/bioptim/models/biorbd/viewer_utils.py +++ b/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: diff --git a/tests/shard3/test_global_torque_driven_ocp.py b/tests/shard3/test_global_torque_driven_ocp.py index 94b8e1f47..78a46f6ba 100644 --- a/tests/shard3/test_global_torque_driven_ocp.py +++ b/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):