fix and refactor for bug free tracked markers

pyomeca · Oct 15, 2024 · 6cfe2ab · 6cfe2ab
1 parent 3c28e69
commit 6cfe2ab
Showing 1 changed file with 13 additions and 13 deletions.
diff --git a/bioptim/models/biorbd/viewer_utils.py b/bioptim/models/biorbd/viewer_utils.py
@@ -19,19 +19,19 @@ def _prepare_tracked_markers_for_animation(
 
         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: