fix: catch a bug in tracked marker animation function

pyomeca · Oct 15, 2024 · 3c28e69 · 3c28e69
1 parent f8a5a1b
commit 3c28e69
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/bioptim/models/biorbd/viewer_utils.py b/bioptim/models/biorbd/viewer_utils.py
@@ -13,7 +13,7 @@ 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
 

diff --git 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):