feat: handling 2d dim with partitionned states

so that we can send timeseries to state_partition
pyomeca · Jan 30, 2025 · 6883210 · 6883210
1 parent bb2cedc
commit 6883210
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/bioptim/models/biorbd/holonomic_biorbd_model.py b/bioptim/models/biorbd/holonomic_biorbd_model.py
@@ -480,9 +480,11 @@ def state_from_partition(self, state_u: MX, state_v: MX) -> MX:
         q = MX() if isinstance(state_u, MX) else DM()
         for i in range(self.nb_q):
             if i in self._independent_joint_index:
-                q = vertcat(q, state_u[self._independent_joint_index.index(i)])
+                slicing = slice(self._independent_joint_index.index(i), self._independent_joint_index.index(i) + 1)
+                q = vertcat(q, state_u[slicing, :])
             else:
-                q = vertcat(q, state_v[self._dependent_joint_index.index(i)])
+                slicing = slice(self._dependent_joint_index.index(i), self._dependent_joint_index.index(i) + 1)
+                q = vertcat(q, state_v[slicing, :])
 
         return q