diff --git a/navipy/trajectories/__init__.py b/navipy/trajectories/__init__.py index 664f8c62f0c3627ff8b5ba27a8ac92b916ed889a..163f3c14c5a9d83c905948d349a1a697cff0f216 100644 --- a/navipy/trajectories/__init__.py +++ b/navipy/trajectories/__init__.py @@ -48,8 +48,14 @@ def _markerstransform(index_i, trajectory, trans_mat = htf.compose_matrix(angles=angles, translate=translate, axes=rotation_mode) - tmarker = trans_mat.dot(homogeneous_markers)[:3] - return tmarker.transpose().flatten() + tmarker = trans_mat.dot(homogeneous_markers) + tmarker = pd.DataFrame(data=tmarker, + index=homogeneous_markers.index, + columns=homogeneous_markers.columns) + # We do not need w + tmarker = tmarker.loc[['x', 'y', 'z'], :].unstack() + tmarker.name = index_i + return tmarker def _invmarkerstransform(index_i, trajectory, @@ -60,8 +66,14 @@ def _invmarkerstransform(index_i, trajectory, trans_mat = htf.compose_matrix(angles=angles, translate=translate, axes=rotation_mode) - tmarker = np.linalg.inv(trans_mat).dot(homogeneous_markers)[:3] - return tmarker.transpose().flatten() + tmarker = np.linalg.inv(trans_mat).dot(homogeneous_markers) + tmarker = pd.DataFrame(data=tmarker, + index=homogeneous_markers.index, + columns=homogeneous_markers.columns) + # We do not need w + tmarker = tmarker.loc[['x', 'y', 'z'], :].unstack() + tmarker.name = index_i + return tmarker class Trajectory(pd.DataFrame): @@ -511,6 +523,11 @@ class Trajectory(pd.DataFrame): # The marker are assume to be a multiIndex dataframe homogeneous_markers = markers.unstack() homogeneous_markers['w'] = 1 + # Make sure that columns are correctly ordered + homogeneous_markers = homogeneous_markers[['x', 'y', 'z', 'w']] + # Transpose because we apply homogeneous transformation + # on the marker, and thus a 4x4 matrix on a 4xN matrix + # here N is the number of markers homogeneous_markers = homogeneous_markers.transpose() # Looping throught the indeces # to get the homogeneous transformation from the position orientation @@ -522,6 +539,8 @@ class Trajectory(pd.DataFrame): homogeneous_markers=homogeneous_markers, rotation_mode=self.rotation_mode), indeces) + # unwrap results + indeces = [res.name for res in result] transformed_markers = pd.DataFrame(data=result, index=indeces, columns=markers.index, @@ -551,6 +570,11 @@ class Trajectory(pd.DataFrame): # The marker are assume to be a multiIndex dataframe homogeneous_markers = markers.unstack() homogeneous_markers['w'] = 1 + # Make sure that columns are correctly ordered + homogeneous_markers = homogeneous_markers[['x', 'y', 'z', 'w']] + # Transpose because we apply homogeneous transformation + # on the marker, and thus a 4x4 matrix on a 4xN matrix + # here N is the number of markers homogeneous_markers = homogeneous_markers.transpose() # Looping throught the indeces # to get the homogeneous transformation from the position orientation @@ -562,6 +586,8 @@ class Trajectory(pd.DataFrame): homogeneous_markers=homogeneous_markers, rotation_mode=self.rotation_mode), indeces) + # unwrap results + indeces = [res.name for res in result] transformed_markers = pd.DataFrame(data=result, index=indeces, columns=markers.index,