Speed up parent link indexes extraction

ami-iit · Oct 23, 2024 · 4c631fa · 4c631fa
1 parent c8d6b53
commit 4c631fa
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/jaxsim/api/references.py b/src/jaxsim/api/references.py
@@ -492,9 +492,9 @@ def apply_frame_forces(
 
         # Extract the frame indices.
         frame_idxs = js.frame.names_to_idxs(frame_names=frame_names, model=model)
-        parent_link_idxs = jax.vmap(js.frame.idx_of_parent_link, in_axes=(None,))(
-            model, frame_index=frame_idxs
-        )
+        parent_link_idxs = model.kin_dyn_parameters.frame_parameters.body[
+            frame_idxs - model.number_of_links
+        ]
 
         exceptions.raise_value_error_if(
             condition=~data.valid(model=model),